[Model] add GLM-5.3-Flash support (#53906)
Jiangyun Zhu committed
Sep 4, 2026 at 00:40 UTC
98ed0856f31fa3aaf5e27464e2b4ef5a8ee6b2f5
94 files changed
+17926
-265
.buildkite/test_areas/kernels.yaml
+6
-1
@@ -166,7 +166,7 @@ steps:
166
commands:
167
- pytest -v -s kernels/attention/test_triton_unified_attention_diffkv.py
168
169
-- label: ":nvidia: (H100) FlashMLA Kernels"
169
+- label: ":nvidia: (H100) MLA Kernel Test"
170
key: kernels-flashmla-test-h100
171
timeout_in_minutes: 25
172
device: h100
@@ -176,12 +176,16 @@ steps:
176
- vllm/v1/attention/ops/flashmla.py
177
- vllm/v1/attention/backends/mla/flashmla.py
178
- vllm/v1/attention/backends/mla/flashmla_sparse.py
179
+ - vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py
180
+ - vllm/utils/flashinfer.py
181
+ - tests/kernels/attention/test_flashinfer_mla_decode.py
182
- tests/kernels/attention/test_flashmla.py
183
- tests/kernels/attention/test_flashmla_sparse.py
184
- tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py
185
commands:
186
- pytest -v -s kernels/attention/test_flashmla.py
187
- pytest -v -s kernels/attention/test_flashmla_sparse.py
188
+ - pytest -v -s kernels/attention/test_flashinfer_mla_decode.py
189
- pytest -v -s kernels/attention/test_mla_cross_layer_kernel_equivalence.py
190
191
- label: ":nvidia: (L4) Quantization Kernels Shard %N"
@@ -337,6 +341,7 @@ steps:
341
- csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu
342
- csrc/libtorch_stable/ops.h
343
- csrc/libtorch_stable/torch_bindings.cpp
344
+ - tests/kernels/attention/test_flashinfer_mla_decode.py
345
- tests/kernels/attention/test_minimax_m3_msa_cutlass_sparse_decode.py
346
- tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py
347
- tests/kernels/mamba/test_gdn_prefill_cutedsl.py
csrc/libtorch_stable/cache_kernels.cu
+9
-4
@@ -1268,6 +1268,9 @@ __global__ void gather_and_maybe_dequant_cache_page(
1268
#define CALL_GATHER_CACHE_576(SCALAR_T, CACHE_T, KV_DTYPE) \
1269
CALL_GATHER_CACHE(SCALAR_T, CACHE_T, KV_DTYPE, 576)
1270
1271
+#define CALL_GATHER_CACHE_512(SCALAR_T, CACHE_T, KV_DTYPE) \
1272
+ CALL_GATHER_CACHE(SCALAR_T, CACHE_T, KV_DTYPE, 512)
1273
+
1274
#define CALL_GATHER_CACHE_320(SCALAR_T, CACHE_T, KV_DTYPE) \
1275
CALL_GATHER_CACHE(SCALAR_T, CACHE_T, KV_DTYPE, 320)
1276
@@ -1305,10 +1308,9 @@ void gather_and_maybe_dequant_cache(
1308
seq_starts.value().scalar_type() == torch::headeronly::ScalarType::Int,
1309
"seq_starts must be int32");
1310
}
1308
- STD_TORCH_CHECK(
1309
- head_dim == 320 || head_dim == 576,
1310
- "gather_and_maybe_dequant_cache only support the head_dim to 320 or 576 "
1311
- "for better performance")
1311
+ STD_TORCH_CHECK(head_dim == 320 || head_dim == 512 || head_dim == 576,
1312
+ "gather_and_maybe_dequant_cache only support the head_dim to "
1313
+ "320 or 512 or 576 for better performance")
1314
1315
STD_TORCH_CHECK(src_cache.device() == dst.device(),
1316
"src_cache and dst must be on the same device");
@@ -1344,6 +1346,9 @@ void gather_and_maybe_dequant_cache(
1346
if (head_dim == 576) {
1347
DISPATCH_BY_KV_CACHE_DTYPE(dst.scalar_type(), kv_cache_dtype,
1348
CALL_GATHER_CACHE_576);
1349
+ } else if (head_dim == 512) {
1350
+ DISPATCH_BY_KV_CACHE_DTYPE(dst.scalar_type(), kv_cache_dtype,
1351
+ CALL_GATHER_CACHE_512);
1352
} else {
1353
DISPATCH_BY_KV_CACHE_DTYPE(dst.scalar_type(), kv_cache_dtype,
1354
CALL_GATHER_CACHE_320);
tests/kernels/attention/test_flashinfer_mla_decode.py
+168
-2
@@ -9,13 +9,23 @@ from vllm.platforms import current_platform
9
10
FLASHINFER_WORKSPACE_BUFFER_SIZE = 128 * 1024 * 1024
11
12
-if not current_platform.has_device_capability(100):
12
+if not current_platform.is_cuda() or not current_platform.has_device_capability(90):
13
pytest.skip(
14
- reason="FlashInfer MLA Requires compute capability of 10 or above.",
14
+ reason="FlashInfer MLA requires CUDA compute capability 9.0 or above.",
15
allow_module_level=True,
16
)
17
else:
18
from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
19
+ from flashinfer.mla import BatchMLAPagedAttentionWrapper
20
+
21
+requires_sm90 = pytest.mark.skipif(
22
+ not current_platform.is_device_capability_family(90),
23
+ reason="This test requires an SM90 GPU.",
24
+)
25
+requires_sm10x = pytest.mark.skipif(
26
+ not current_platform.is_device_capability_family(100),
27
+ reason="This test requires an SM10x GPU.",
28
+)
29
30
# Deepseek R1 MLA config.
31
NUM_HEADS = 128
@@ -82,6 +92,7 @@ def ref_mla(
92
@pytest.mark.parametrize("dtype", [torch.bfloat16])
93
@pytest.mark.parametrize("bs", [1, 2, 4, 16])
94
@pytest.mark.parametrize("block_size", [32, 64])
95
+@requires_sm10x
96
def test_flashinfer_mla_decode(dtype: torch.dtype, bs: int, block_size: int):
97
torch.set_default_device("cuda")
98
torch.manual_seed(42)
@@ -119,6 +130,161 @@ def test_flashinfer_mla_decode(dtype: torch.dtype, bs: int, block_size: int):
130
torch.testing.assert_close(out_ans, out_ref, atol=1e-2, rtol=1e-2)
131
132
133
+@requires_sm10x
134
+def test_flashinfer_trtllm_sparse_mla_decode_without_rope():
135
+ """The native sparse MLA path supports a zero-width rotary tail."""
136
+ torch.set_default_device("cuda")
137
+ torch.manual_seed(42)
138
+
139
+ batch_size = 2
140
+ block_size = 64
141
+ num_blocks = 4
142
+ sparse_topk = 128
143
+ valid_lens = torch.tensor([17, 73], dtype=torch.int32)
144
+
145
+ query = torch.randn(
146
+ batch_size,
147
+ 1,
148
+ NUM_HEADS,
149
+ KV_LORA_RANK,
150
+ dtype=torch.bfloat16,
151
+ )
152
+ kv_cache = torch.randn(
153
+ num_blocks,
154
+ block_size,
155
+ KV_LORA_RANK,
156
+ dtype=torch.bfloat16,
157
+ )
158
+
159
+ num_slots = num_blocks * block_size
160
+ slot_tables = torch.stack(
161
+ [torch.randperm(num_slots)[:sparse_topk] for _ in range(batch_size)]
162
+ ).to(torch.int32)
163
+ for row, valid_len in zip(slot_tables, valid_lens.tolist()):
164
+ row[valid_len:] = -1
165
+
166
+ workspace_buffer = torch.empty(
167
+ FLASHINFER_WORKSPACE_BUFFER_SIZE,
168
+ dtype=torch.int8,
169
+ )
170
+ out = trtllm_batch_decode_with_kv_cache_mla(
171
+ query=query,
172
+ kv_cache=kv_cache.unsqueeze(1),
173
+ workspace_buffer=workspace_buffer,
174
+ qk_nope_head_dim=QK_NOPE_HEAD_DIM,
175
+ kv_lora_rank=KV_LORA_RANK,
176
+ qk_rope_head_dim=0,
177
+ block_tables=slot_tables.unsqueeze(1),
178
+ seq_lens=valid_lens,
179
+ max_seq_len=sparse_topk,
180
+ sparse_mla_top_k=sparse_topk,
181
+ sparse_mla_top_k_lens=valid_lens,
182
+ bmm1_scale=QK_NOPE_HEAD_DIM**-0.5,
183
+ bmm2_scale=1.0,
184
+ ).squeeze(1)
185
+
186
+ flat_cache = kv_cache.view(num_slots, KV_LORA_RANK).float()
187
+ refs = []
188
+ for batch_idx, valid_len in enumerate(valid_lens.tolist()):
189
+ selected_kv = flat_cache[slot_tables[batch_idx, :valid_len].long()]
190
+ scores = torch.einsum("hd,kd->hk", query[batch_idx, 0].float(), selected_kv)
191
+ probs = torch.softmax(scores * QK_NOPE_HEAD_DIM**-0.5, dim=-1)
192
+ refs.append(torch.einsum("hk,kd->hd", probs, selected_kv))
193
+ ref = torch.stack(refs).to(torch.bfloat16)
194
+
195
+ torch.testing.assert_close(out, ref, atol=2e-2, rtol=2e-2)
196
+
197
+
198
+@requires_sm90
199
+def test_flashinfer_sm90_fp8_mla_decode_without_rope():
200
+ """Hopper FA3 supports BF16 queries over an FP8 cache without KPE."""
201
+ torch.manual_seed(42)
202
+ device = torch.device("cuda")
203
+ batch_size = 2
204
+ num_heads = 16
205
+ page_size = 16
206
+ num_pages = 6
207
+
208
+ q_nope = torch.randn(
209
+ batch_size,
210
+ num_heads,
211
+ KV_LORA_RANK,
212
+ dtype=torch.bfloat16,
213
+ device=device,
214
+ )
215
+ q_pe = torch.empty(
216
+ batch_size,
217
+ num_heads,
218
+ 0,
219
+ dtype=torch.bfloat16,
220
+ device=device,
221
+ )
222
+
223
+ ckv = torch.randn(
224
+ num_pages,
225
+ page_size,
226
+ KV_LORA_RANK,
227
+ device=device,
228
+ )
229
+ fp8_max = torch.finfo(torch.float8_e4m3fn).max
230
+ ckv_scale = ckv.abs().max().item() / fp8_max
231
+ ckv_fp8 = (ckv / ckv_scale).clamp(-fp8_max, fp8_max).to(torch.float8_e4m3fn)
232
+ scale_bf16 = torch.tensor(ckv_scale, dtype=torch.bfloat16, device=device)
233
+ ckv_ref = ckv_fp8.to(torch.bfloat16) * scale_bf16
234
+ kpe_fp8 = torch.empty(
235
+ num_pages,
236
+ page_size,
237
+ 0,
238
+ dtype=torch.float8_e4m3fn,
239
+ device=device,
240
+ )
241
+ kpe_ref = torch.empty(
242
+ num_pages,
243
+ page_size,
244
+ 0,
245
+ dtype=torch.bfloat16,
246
+ device=device,
247
+ )
248
+
249
+ qo_indptr = torch.tensor([0, 1, 2], dtype=torch.int32, device=device)
250
+ kv_indptr = torch.tensor([0, 3, 5], dtype=torch.int32, device=device)
251
+ kv_indices = torch.tensor([4, 1, 3, 0, 5], dtype=torch.int32, device=device)
252
+ kv_lens = torch.tensor([45, 29], dtype=torch.int32, device=device)
253
+ sm_scale = QK_NOPE_HEAD_DIM**-0.5
254
+
255
+ def run(
256
+ ckv_cache: torch.Tensor,
257
+ kpe_cache: torch.Tensor,
258
+ **kwargs,
259
+ ) -> torch.Tensor:
260
+ workspace = torch.empty(
261
+ FLASHINFER_WORKSPACE_BUFFER_SIZE,
262
+ dtype=torch.uint8,
263
+ device=device,
264
+ )
265
+ wrapper = BatchMLAPagedAttentionWrapper(workspace, backend="fa3")
266
+ wrapper.plan(
267
+ qo_indptr,
268
+ kv_indptr,
269
+ kv_indices,
270
+ kv_lens,
271
+ num_heads,
272
+ KV_LORA_RANK,
273
+ 0,
274
+ page_size,
275
+ False,
276
+ sm_scale,
277
+ q_data_type=torch.bfloat16,
278
+ kv_data_type=ckv_cache.dtype,
279
+ )
280
+ return wrapper.run(q_nope, q_pe, ckv_cache, kpe_cache, **kwargs)
281
+
282
+ out_ref = run(ckv_ref, kpe_ref)
283
+ out = run(ckv_fp8, kpe_fp8, ckv_scale=ckv_scale, kpe_scale=1.0)
284
+ torch.testing.assert_close(out, out_ref, atol=2e-2, rtol=2e-2)
285
+
286
+
287
+@requires_sm10x
288
def test_flashinfer_mla_decode_workspace_supports_autotune():
289
"""vLLM's FlashInfer MLA decode workspace must be int8 for autotuning.
290
tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py
+2
@@ -41,6 +41,8 @@ def _make_builder():
41
builder.device = torch.device("cpu")
42
builder.kv_cache_spec = SimpleNamespace(block_size=1)
43
builder.model_dtype = torch.bfloat16
44
+ builder.kv_cache_dtype = "fp8"
45
+ builder.mla_dims = SimpleNamespace(kv_lora_rank=512, qk_rope_head_dim=64)
46
builder.topk_tokens = topk_tokens
47
builder.req_id_per_token_buffer = torch.zeros(
48
max_num_batched_tokens, dtype=torch.int32, device="cpu"
tests/kernels/mamba/test_gdn_prefill_flashinfer.py
new
+53
@@ -0,0 +1,53 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+
7
+from vllm.platforms import current_platform
8
+
9
+if current_platform.is_rocm():
10
+ pytest.skip(
11
+ reason="FlashInfer GDN prefill is not supported on ROCm.",
12
+ allow_module_level=True,
13
+ )
14
+
15
+import flashinfer.gdn_prefill # noqa: E402
16
+
17
+from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (
18
+ fi_chunk_gated_delta_rule,
19
+) # noqa: E402
20
+
21
+
22
+def test_flashinfer_gdn_prefill_uses_int64_cu_seqlens(monkeypatch):
23
+ captured_cu_seqlens = None
24
+
25
+ def fake_chunk_gated_delta_rule(**kwargs):
26
+ nonlocal captured_cu_seqlens
27
+ captured_cu_seqlens = kwargs["cu_seqlens"]
28
+ return kwargs["q"]
29
+
30
+ monkeypatch.setattr(
31
+ flashinfer.gdn_prefill,
32
+ "chunk_gated_delta_rule",
33
+ fake_chunk_gated_delta_rule,
34
+ )
35
+ q = torch.zeros(1, 2, 1, 2)
36
+ cu_seqlens = torch.tensor([0, 2], dtype=torch.int32)
37
+
38
+ output, final_state = fi_chunk_gated_delta_rule(
39
+ q=q,
40
+ k=q,
41
+ v=q,
42
+ g=torch.zeros(1, 2, 1),
43
+ beta=torch.zeros(1, 2, 1),
44
+ initial_state=torch.zeros(1, 1, 2, 2),
45
+ output_final_state=False,
46
+ cu_seqlens=cu_seqlens,
47
+ use_qk_l2norm_in_kernel=False,
48
+ )
49
+
50
+ assert captured_cu_seqlens is not None
51
+ assert captured_cu_seqlens.dtype == torch.int64
52
+ assert output.shape == q.shape
53
+ assert final_state is None
tests/kernels/test_kpool_decode_update_batched.py
new
+549
@@ -0,0 +1,549 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Unit test for the batched kpool decode-update kernel.
4
+
5
+Validates ``kpool_decode_update_and_maybe_write_cache_batched`` against an
6
+independent pure-torch reference that replicates the per-request, in-position
7
+order semantics: stash each token into a paged tail ring; on pool completion
8
+(``pos % pool_size == pool_size-1``) softmax(gate+ape)-weighted sum + Hadamard-128
9
++ per-vector fp8 absmax quant + write to the indexer K cache. Covers
10
+no-completion, completion-at-end, completion-mid-batch, non-uniform padding,
11
+plain decode, plus a randomized fuzz pass.
12
+
13
+The kernel iterates each request's ``next_n`` tokens in position order inside
14
+one program (grid = num_requests) to preserve the pool-completion
15
+read-after-stash dependency; the reference mirrors that ordering.
16
+
17
+``test_decode_writer_matches_prefill_writer`` is deliberately NOT
18
+reference-based: it checks the decode writer against the *prefill* writer
19
+(``kpool_compress_and_write_cache``), the invariant that actually matters in
20
+production. A hand-written reference can drift to match a buggy kernel -- that
21
+is exactly how the stash-gating bug (intra-pool tokens never entering the tail
22
+ring, because the stash was gated on the pool-granular ``slot_mapping``) stayed
23
+green here.
24
+"""
25
+
26
+import math
27
+
28
+import pytest
29
+import torch
30
+
31
+from vllm.platforms import current_platform
32
+
33
+if current_platform.is_rocm():
34
+ from vllm.models.glm5next.amd.ops.kpool_compress import (
35
+ kpool_compress_and_write_cache,
36
+ kpool_decode_update_and_maybe_write_cache_batched,
37
+ kpool_seed_tail_cache,
38
+ )
39
+else:
40
+ from vllm.models.glm5next.nvidia.ops.kpool_compress import (
41
+ kpool_compress_and_write_cache,
42
+ kpool_decode_update_and_maybe_write_cache_batched,
43
+ kpool_seed_tail_cache,
44
+ )
45
+
46
+HEAD_DIM = 128
47
+POOL_SIZE = 16
48
+PAGE_SIZE = 64
49
+NUM_BLOCKS = 32
50
+ROUND_SCALE = True
51
+FP8_DTYPE = current_platform.fp8_dtype()
52
+FP8_MAX = torch.finfo(FP8_DTYPE).max
53
+
54
+
55
+def _make_caches():
56
+ kv = torch.zeros(
57
+ NUM_BLOCKS, PAGE_SIZE, HEAD_DIM + 4, dtype=torch.uint8, device="cuda"
58
+ )
59
+ tail = torch.zeros(
60
+ NUM_BLOCKS, 2, POOL_SIZE, HEAD_DIM, dtype=torch.bfloat16, device="cuda"
61
+ )
62
+ return kv, tail
63
+
64
+
65
+def _tail_slot_for(blocks, pos):
66
+ """tail_slot = block*POOL + pos%POOL; each request owns a distinct tail block."""
67
+ blk = torch.tensor(blocks, device=pos.device, dtype=torch.int32).unsqueeze(1)
68
+ return (blk * POOL_SIZE + pos % POOL_SIZE).to(torch.int32)
69
+
70
+
71
+def _seed_prior(tail, blocks, n_prior, seed=42):
72
+ if n_prior <= 0:
73
+ return
74
+ g = torch.Generator(device=tail.device).manual_seed(seed)
75
+ prior_k = torch.randn(
76
+ len(blocks),
77
+ n_prior,
78
+ HEAD_DIM,
79
+ dtype=torch.bfloat16,
80
+ device=tail.device,
81
+ generator=g,
82
+ )
83
+ prior_s = torch.randn(
84
+ len(blocks),
85
+ n_prior,
86
+ HEAD_DIM,
87
+ dtype=torch.bfloat16,
88
+ device=tail.device,
89
+ generator=g,
90
+ )
91
+ for i, blk in enumerate(blocks):
92
+ tail[blk, 0, :n_prior, :] = prior_k[i]
93
+ tail[blk, 1, :n_prior, :] = prior_s[i]
94
+
95
+
96
+def _hadamard128_torch(x: torch.Tensor) -> torch.Tensor:
97
+ """Reference Hadamard-128 on the last dim (must be 128)."""
98
+ n = x.shape[-1]
99
+ assert n == 128
100
+ h = torch.tensor([[1.0, 1.0], [1.0, -1.0]], dtype=torch.float32, device=x.device)
101
+ while h.shape[0] < n:
102
+ h = torch.cat([torch.cat([h, h], dim=1), torch.cat([h, -h], dim=1)], dim=0)
103
+ h = h / math.sqrt(n)
104
+ return x @ h
105
+
106
+
107
+def _torch_reference(
108
+ kv: torch.Tensor,
109
+ tail: torch.Tensor,
110
+ tail_slot: torch.Tensor,
111
+ key: torch.Tensor,
112
+ score: torch.Tensor,
113
+ ape: torch.Tensor,
114
+ slot_map: torch.Tensor,
115
+ pos: torch.Tensor,
116
+) -> tuple[torch.Tensor, torch.Tensor]:
117
+ """Independent reference for the batched decode-update kernel.
118
+
119
+ For each request, iterate its next_n tokens in order. On each token:
120
+ - if pos%POOL == POOL-1 and pos_valid: compress the pool (slots
121
+ [pool_start..pool_start+POOL-1], current token via is_current) and write
122
+ fp8 K + fp32 scale to kv_cache at cache_loc.
123
+ - always stash the current token's K/score into tail[block, pos%POOL].
124
+ """
125
+ kv = kv.clone()
126
+ tail = tail.clone()
127
+ B, next_n = pos.shape
128
+ # The indexer K cache is [num_blocks, PAGE_SIZE, HEAD_DIM+4] uint8 but the
129
+ # kernels interpret each page as [HEAD_DIM*PAGE_SIZE bytes of K (token-major)
130
+ # | 4*PAGE_SIZE bytes of fp32 scale (token-major)]. Operate on a flat byte
131
+ # view so the reference writes K and scale at the exact offsets the kernel
132
+ # uses (page_base + tok*HEAD_DIM for K; page_base + HEAD_DIM*PAGE_SIZE +
133
+ # tok*4 for scale).
134
+ page_bytes = PAGE_SIZE * (HEAD_DIM + 4)
135
+ k_region = HEAD_DIM * PAGE_SIZE
136
+ tail_slot_cpu = tail_slot.cpu().tolist()
137
+ slot_map_cpu = slot_map.cpu().tolist()
138
+ pos_cpu = pos.cpu().tolist()
139
+ key_cpu = key.float().cpu()
140
+ score_cpu = score.float().cpu()
141
+ ape_cpu = ape.cpu()
142
+ tail_cpu = tail.float().cpu()
143
+ kv_flat = kv.view(torch.uint8).reshape(-1).cpu()
144
+
145
+ for b in range(B):
146
+ for t in range(next_n):
147
+ cache_loc = slot_map_cpu[b][t]
148
+ p = pos_cpu[b][t]
149
+ pos_valid = cache_loc >= 0 and p >= 0
150
+ safe_pos = max(p, 0)
151
+ slot = safe_pos % POOL_SIZE
152
+ phys_slot = safe_pos % POOL_SIZE
153
+ # Per-token block derivation (a leading invalid sentinel must not
154
+ # poison the base for the rest of the request); clamped like the
155
+ # kernel so an invalid entry can't form a negative base.
156
+ block = max(tail_slot_cpu[b][t], 0) // POOL_SIZE
157
+
158
+ cur_key = key_cpu[b, t]
159
+ cur_score = score_cpu[b, t]
160
+
161
+ if pos_valid and slot == POOL_SIZE - 1:
162
+ pool_logical_start = safe_pos - slot
163
+ pool_scores = []
164
+ pool_ks = []
165
+ for ps in range(POOL_SIZE):
166
+ is_current = ps == slot
167
+ phys = (pool_logical_start + ps) % POOL_SIZE
168
+ if is_current:
169
+ s = cur_score
170
+ k = cur_key
171
+ else:
172
+ s = tail_cpu[block, 1, phys]
173
+ k = tail_cpu[block, 0, phys]
174
+ s = s + ape_cpu[ps]
175
+ pool_scores.append(s)
176
+ pool_ks.append(k)
177
+ pool_scores = torch.stack(pool_scores) # [POOL, D]
178
+ pool_ks = torch.stack(pool_ks) # [POOL, D]
179
+ max_score = pool_scores.max(dim=0).values
180
+ prob = torch.exp(pool_scores - max_score)
181
+ denom = prob.sum(dim=0)
182
+ acc = (pool_ks * prob).sum(dim=0)
183
+ x = (acc / denom).to(torch.bfloat16).to(torch.float32)
184
+ x = _hadamard128_torch(x).to(torch.bfloat16).to(torch.float32)
185
+ absmax = torch.clamp(x.abs().max(), min=1e-4)
186
+ if ROUND_SCALE:
187
+ scale = torch.exp2(torch.ceil(torch.log2(absmax / FP8_MAX)))
188
+ else:
189
+ scale = absmax / FP8_MAX
190
+ quantized = torch.clamp(x / scale, -FP8_MAX, FP8_MAX).to(FP8_DTYPE)
191
+ # write K and scale at the separated-layout offsets
192
+ loc = cache_loc
193
+ loc_page_index = loc // PAGE_SIZE
194
+ loc_tok = loc % PAGE_SIZE
195
+ page_base = loc_page_index * page_bytes
196
+ if current_platform.is_rocm():
197
+ dims = torch.arange(HEAD_DIM)
198
+ k_off = (
199
+ page_base
200
+ + (loc_tok // 16) * 16 * HEAD_DIM
201
+ + (dims // 16) * 16 * 16
202
+ + (loc_tok % 16) * 16
203
+ + dims % 16
204
+ )
205
+ else:
206
+ k_off = page_base + loc_tok * HEAD_DIM + torch.arange(HEAD_DIM)
207
+ s_off = page_base + k_region + loc_tok * 4
208
+ kv_flat[k_off] = quantized.view(torch.uint8)
209
+ kv_flat[s_off : s_off + 4] = scale.detach().reshape(1).view(torch.uint8)
210
+
211
+ # stash -- gated on the TOKEN-granular tail slot, not on pos_valid.
212
+ # pos_valid keys off the pool-granular cache_loc, which is -1 for
213
+ # every token that is not the pool's last, so gating the stash on it
214
+ # would drop all intra-pool tokens.
215
+ if p >= 0 and tail_slot_cpu[b][t] >= 0:
216
+ tail_cpu[block, 0, phys_slot] = cur_key
217
+ tail_cpu[block, 1, phys_slot] = cur_score
218
+
219
+ kv_out = kv_flat.view(NUM_BLOCKS, PAGE_SIZE, HEAD_DIM + 4).to(device="cuda")
220
+ return kv_out, tail_cpu.to(torch.bfloat16).to(device="cuda")
221
+
222
+
223
+def _assert_eq(r_ref, r_kern):
224
+ kv_ref, tail_ref = r_ref
225
+ kv_kern, tail_kern = r_kern
226
+ assert torch.equal(kv_ref, kv_kern), (
227
+ "kv_cache differs: max diff "
228
+ f"{(kv_ref.int() - kv_kern.int()).abs().max().item()}"
229
+ )
230
+ assert torch.equal(tail_ref, tail_kern), (
231
+ "tail_kv_cache differs: max diff "
232
+ f"{(tail_ref.float() - tail_kern.float()).abs().max().item()}"
233
+ )
234
+
235
+
236
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm required")
237
+def test_amd_prefill_writer_uses_preshuffled_cache_layout():
238
+ from vllm.models.glm5next.amd.ops.kpool_compress import (
239
+ kpool_compress_and_write_cache as amd_kpool_compress,
240
+ )
241
+
242
+ torch.manual_seed(0)
243
+ token_offset = 17
244
+ kv = torch.zeros(1, PAGE_SIZE, HEAD_DIM + 4, dtype=torch.uint8, device="cuda")
245
+ key = torch.randn(1, POOL_SIZE, HEAD_DIM, dtype=torch.bfloat16, device="cuda")
246
+ score = torch.randn_like(key)
247
+ ape = torch.randn(POOL_SIZE, HEAD_DIM, dtype=torch.float32, device="cuda")
248
+ compressed_k, compressed_scale = amd_kpool_compress(
249
+ kv,
250
+ key,
251
+ score,
252
+ ape,
253
+ torch.tensor([token_offset], dtype=torch.int64, device="cuda"),
254
+ pool_size=POOL_SIZE,
255
+ head_dim=HEAD_DIM,
256
+ round_scale=ROUND_SCALE,
257
+ return_compressed=True,
258
+ )
259
+
260
+ dim = torch.arange(HEAD_DIM, device="cuda")
261
+ offsets = (
262
+ (token_offset // 16) * 16 * HEAD_DIM
263
+ + (dim // 16) * 16 * 16
264
+ + (token_offset % 16) * 16
265
+ + dim % 16
266
+ )
267
+ flat = kv.view(torch.uint8).reshape(-1)
268
+ stored_k = flat[offsets].view(compressed_k.dtype)
269
+ scale_offset = PAGE_SIZE * HEAD_DIM + token_offset * 4
270
+ stored_scale = flat[scale_offset : scale_offset + 4].view(torch.float32)
271
+
272
+ assert torch.equal(stored_k, compressed_k[0])
273
+ assert torch.equal(stored_scale, compressed_scale)
274
+
275
+
276
+def _run_kernel(kv, tail, tail_slot, key, score, ape, slot_map, pos):
277
+ kv = kv.clone()
278
+ tail = tail.clone()
279
+ kpool_decode_update_and_maybe_write_cache_batched(
280
+ kv,
281
+ tail,
282
+ tail_slot,
283
+ key,
284
+ score,
285
+ ape,
286
+ slot_map,
287
+ pos,
288
+ POOL_SIZE,
289
+ HEAD_DIM,
290
+ round_scale=ROUND_SCALE,
291
+ )
292
+ return kv, tail
293
+
294
+
295
+@pytest.mark.parametrize("pool_size", [4, 16])
296
+def test_decode_writer_matches_prefill_writer(pool_size):
297
+ """Compare production decode and prefill writers for pool sizes 4 and 16."""
298
+ n_pools, page, nblk = 8, 64, 4
299
+ n_tok = n_pools * pool_size
300
+ dev = "cuda"
301
+ torch.manual_seed(0)
302
+ k = torch.randn(n_tok, HEAD_DIM, dtype=torch.bfloat16, device=dev)
303
+ score = torch.randn(n_tok, HEAD_DIM, dtype=torch.bfloat16, device=dev)
304
+ ape = torch.randn(pool_size, HEAD_DIM, dtype=torch.float32, device=dev)
305
+
306
+ kv_prefill = torch.zeros(nblk, page, HEAD_DIM + 4, dtype=torch.uint8, device=dev)
307
+ kpool_compress_and_write_cache(
308
+ kv_prefill,
309
+ k.view(n_pools, pool_size, HEAD_DIM),
310
+ score.view(n_pools, pool_size, HEAD_DIM),
311
+ ape,
312
+ torch.arange(n_pools, dtype=torch.int64, device=dev),
313
+ pool_size=pool_size,
314
+ head_dim=HEAD_DIM,
315
+ round_scale=ROUND_SCALE,
316
+ )
317
+
318
+ # One request owning tail block 0, fed one token per decode step.
319
+ kv_decode = torch.zeros_like(kv_prefill)
320
+ tail = torch.zeros(nblk, 2, pool_size, HEAD_DIM, dtype=torch.bfloat16, device=dev)
321
+ for t in range(n_tok):
322
+ completes = t % pool_size == pool_size - 1
323
+ kpool_decode_update_and_maybe_write_cache_batched(
324
+ kv_decode,
325
+ tail,
326
+ # token-granular: every token has a valid tail slot
327
+ torch.tensor([[t % pool_size]], dtype=torch.int32, device=dev),
328
+ k[t].view(1, 1, HEAD_DIM),
329
+ score[t].view(1, 1, HEAD_DIM),
330
+ ape,
331
+ # pool-granular: only the pool's last token carries a cache slot
332
+ torch.tensor(
333
+ [[t // pool_size if completes else -1]], dtype=torch.int32, device=dev
334
+ ),
335
+ torch.tensor([[t]], dtype=torch.int32, device=dev),
336
+ pool_size,
337
+ HEAD_DIM,
338
+ round_scale=ROUND_SCALE,
339
+ )
340
+
341
+ differing = [
342
+ p
343
+ for p in range(n_pools)
344
+ if not torch.equal(
345
+ kv_prefill[p // page, p % page], kv_decode[p // page, p % page]
346
+ )
347
+ ]
348
+ assert not differing, (
349
+ f"decode-written pools differ from prefill-written pools: "
350
+ f"{len(differing)}/{n_pools} (pool_size={pool_size}, first={differing[:5]})"
351
+ )
352
+
353
+
354
+def test_leading_invalid_tail_slot():
355
+ """A request whose FIRST token carries an invalid (-1) tail slot while a
356
+ later token is a real pool completion.
357
+
358
+ The tail block must be derived per token, not from token 0: a leading
359
+ invalid sentinel would otherwise poison the base address for the whole
360
+ request (out-of-bounds tail reads on the completion).
361
+ """
362
+ torch.manual_seed(0)
363
+ B, next_n, blocks = 2, 4, [3, 5]
364
+ # req 0: token 0 invalid (pos -1), tokens 1..3 valid, completion at pos 15
365
+ # req 1: all valid, no completion
366
+ pos = torch.tensor(
367
+ [[-1, 13, 14, 15], [4, 5, 6, 7]], dtype=torch.int32, device="cuda"
368
+ )
369
+ safe_pos = torch.where(pos >= 0, pos, 0)
370
+ tail_slot = _tail_slot_for(blocks, safe_pos)
371
+ # leading invalid entry carries the -1 sentinel, as the scatter path emits
372
+ tail_slot[0, 0] = -1
373
+ slot_map = torch.full((B, next_n), -1, dtype=torch.int32, device="cuda")
374
+ slot_map[0, 3] = 15 # req 0 completes its pool on the last verify token
375
+
376
+ key = torch.randn(B, next_n, HEAD_DIM, dtype=torch.bfloat16, device="cuda")
377
+ score = torch.randn(B, next_n, HEAD_DIM, dtype=torch.bfloat16, device="cuda")
378
+ ape = torch.randn(POOL_SIZE, HEAD_DIM, dtype=torch.float32, device="cuda")
379
+
380
+ kv, tail = _make_caches()
381
+ _seed_prior(tail, blocks, 13)
382
+ r_ref = _torch_reference(kv, tail, tail_slot, key, score, ape, slot_map, pos)
383
+ r_kern = _run_kernel(kv, tail, tail_slot, key, score, ape, slot_map, pos)
384
+ _assert_eq(r_ref, r_kern)
385
+
386
+
387
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm required")
388
+def test_amd_prefill_seed_honors_padded_tail_block_stride():
389
+ """The tail shares a padded indexer allocation in production."""
390
+ kpool = 4
391
+ num_blocks = 6
392
+ logical_block_elems = 2 * kpool * HEAD_DIM
393
+ padded_block_elems = logical_block_elems + 256
394
+ sentinel = -123.0
395
+ backing = torch.full(
396
+ (num_blocks * padded_block_elems,),
397
+ sentinel,
398
+ dtype=torch.bfloat16,
399
+ device="cuda",
400
+ )
401
+ tail = torch.as_strided(
402
+ backing,
403
+ size=(num_blocks, 2, kpool, HEAD_DIM),
404
+ stride=(padded_block_elems, kpool * HEAD_DIM, HEAD_DIM, 1),
405
+ )
406
+
407
+ block = 3
408
+ ring_slot = 2
409
+ key = torch.arange(HEAD_DIM, dtype=torch.bfloat16, device="cuda").unsqueeze(0)
410
+ score = (key + 256).to(torch.bfloat16)
411
+ tail_slot = torch.tensor(
412
+ [block * kpool + ring_slot], dtype=torch.int32, device="cuda"
413
+ )
414
+
415
+ kpool_seed_tail_cache(tail, key, score, tail_slot, kpool, HEAD_DIM)
416
+ torch.accelerator.synchronize()
417
+
418
+ assert torch.equal(tail[block, 0, ring_slot], key[0])
419
+ assert torch.equal(tail[block, 1, ring_slot], score[0])
420
+
421
+ compact_offset = (block * 2 * kpool + ring_slot) * HEAD_DIM
422
+ assert torch.all(backing[compact_offset : compact_offset + HEAD_DIM] == sentinel)
423
+
424
+
425
+@pytest.mark.parametrize(
426
+ "case_id",
427
+ [
428
+ "no_completion",
429
+ "completion_at_end",
430
+ "completion_mid_batch",
431
+ "non_uniform_padding",
432
+ "plain_decode",
433
+ ],
434
+)
435
+def test_batched_matches_reference(case_id):
436
+ torch.manual_seed(0)
437
+ if case_id == "no_completion":
438
+ B, next_n, blocks = 3, 4, [0, 1, 2]
439
+ pos = (
440
+ torch.arange(next_n, device="cuda", dtype=torch.int32)
441
+ .unsqueeze(0)
442
+ .expand(B, -1)
443
+ .contiguous()
444
+ )
445
+ slot_map = torch.full((B, next_n), -1, dtype=torch.int32, device="cuda")
446
+ n_prior = 0
447
+ elif case_id == "completion_at_end":
448
+ B, next_n, blocks = 2, 4, [0, 1]
449
+ pos = torch.tensor(
450
+ [[12, 13, 14, 15], [12, 13, 14, 15]], dtype=torch.int32, device="cuda"
451
+ )
452
+ slot_map = torch.full((B, next_n), -1, dtype=torch.int32, device="cuda")
453
+ slot_map[:, 3] = torch.tensor(
454
+ [15, PAGE_SIZE + 15], dtype=torch.int32, device="cuda"
455
+ )
456
+ n_prior = POOL_SIZE - next_n
457
+ elif case_id == "completion_mid_batch":
458
+ B, next_n, blocks = 3, 4, [0, 1, 2]
459
+ pos = torch.tensor([[13, 14, 15, 16]] * B, dtype=torch.int32, device="cuda")
460
+ slot_map = torch.full((B, next_n), -1, dtype=torch.int32, device="cuda")
461
+ slot_map[:, 2] = torch.tensor(
462
+ [15, PAGE_SIZE + 15, 2 * PAGE_SIZE + 15], dtype=torch.int32, device="cuda"
463
+ )
464
+ n_prior = 13
465
+ elif case_id == "non_uniform_padding":
466
+ B, next_n, blocks = 2, 4, [0, 1]
467
+ pos = torch.tensor(
468
+ [[12, 13, 14, 15], [12, 13, -1, -1]], dtype=torch.int32, device="cuda"
469
+ )
470
+ slot_map = torch.full((B, next_n), -1, dtype=torch.int32, device="cuda")
471
+ slot_map[0, 3] = 15
472
+ n_prior = POOL_SIZE - 4
473
+ else: # plain_decode
474
+ B, next_n, blocks = 4, 1, [0, 1, 2, 3]
475
+ pos = torch.tensor([[5], [6], [7], [8]], dtype=torch.int32, device="cuda")
476
+ slot_map = torch.full((B, next_n), -1, dtype=torch.int32, device="cuda")
477
+ n_prior = 0
478
+
479
+ if case_id == "non_uniform_padding":
480
+ safe_pos = torch.where(pos >= 0, pos, 0)
481
+ tail_slot = torch.where(pos >= 0, _tail_slot_for(blocks, safe_pos), 0)
482
+ else:
483
+ tail_slot = _tail_slot_for(blocks, pos)
484
+
485
+ key = torch.randn(B, next_n, HEAD_DIM, dtype=torch.bfloat16, device="cuda")
486
+ score = torch.randn(B, next_n, HEAD_DIM, dtype=torch.bfloat16, device="cuda")
487
+ ape = torch.randn(POOL_SIZE, HEAD_DIM, dtype=torch.float32, device="cuda")
488
+
489
+ kv, tail = _make_caches()
490
+ _seed_prior(tail, blocks, n_prior)
491
+ r_ref = _torch_reference(kv, tail, tail_slot, key, score, ape, slot_map, pos)
492
+ r_kern = _run_kernel(kv, tail, tail_slot, key, score, ape, slot_map, pos)
493
+ _assert_eq(r_ref, r_kern)
494
+
495
+
496
+@pytest.mark.parametrize("seed", list(range(20)))
497
+def test_batched_matches_reference_fuzz(seed):
498
+ """Random B / next_n / start positions; covers 0, 1, and multi completion."""
499
+ g = torch.Generator(device="cuda").manual_seed(seed)
500
+ B = int(torch.randint(1, 6, (1,), generator=g, device="cuda").item())
501
+ next_n = int(torch.randint(1, 8, (1,), generator=g, device="cuda").item())
502
+ blocks = list(range(B))
503
+
504
+ starts = torch.randint(0, 33, (B,), generator=g, device="cuda", dtype=torch.int32)
505
+ pos = starts.unsqueeze(1) + torch.arange(
506
+ next_n, device="cuda", dtype=torch.int32
507
+ ).unsqueeze(0)
508
+ tail_slot = _tail_slot_for(blocks, pos)
509
+
510
+ is_completion = pos % POOL_SIZE == POOL_SIZE - 1
511
+ blk = torch.tensor(blocks, device="cuda", dtype=torch.int32).unsqueeze(1)
512
+ pool_slot = blk * PAGE_SIZE + (POOL_SIZE - 1)
513
+ slot_map = torch.where(is_completion, pool_slot, torch.full_like(pos, -1))
514
+
515
+ key = torch.randn(
516
+ B, next_n, HEAD_DIM, dtype=torch.bfloat16, device="cuda", generator=g
517
+ )
518
+ score = torch.randn(
519
+ B, next_n, HEAD_DIM, dtype=torch.bfloat16, device="cuda", generator=g
520
+ )
521
+ ape = torch.randn(
522
+ POOL_SIZE, HEAD_DIM, dtype=torch.float32, device="cuda", generator=g
523
+ )
524
+
525
+ kv, tail = _make_caches()
526
+ prior_g = torch.Generator(device="cuda").manual_seed(seed + 1000)
527
+ for b in range(B):
528
+ n_prior = int(starts[b].item()) % POOL_SIZE
529
+ if n_prior > 0:
530
+ pk = torch.randn(
531
+ n_prior,
532
+ HEAD_DIM,
533
+ dtype=torch.bfloat16,
534
+ device="cuda",
535
+ generator=prior_g,
536
+ )
537
+ ps = torch.randn(
538
+ n_prior,
539
+ HEAD_DIM,
540
+ dtype=torch.bfloat16,
541
+ device="cuda",
542
+ generator=prior_g,
543
+ )
544
+ tail[blocks[b], 0, :n_prior, :] = pk
545
+ tail[blocks[b], 1, :n_prior, :] = ps
546
+
547
+ r_ref = _torch_reference(kv, tail, tail_slot, key, score, ape, slot_map, pos)
548
+ r_kern = _run_kernel(kv, tail, tail_slot, key, score, ape, slot_map, pos)
549
+ _assert_eq(r_ref, r_kern)
tests/kernels/test_mhc_kernels.py
+225
-1
@@ -5,13 +5,23 @@ from types import SimpleNamespace
5
import pytest
6
import torch
7
import torch.nn as nn
8
+import torch.nn.functional as F
9
10
import vllm.model_executor.kernels.mhc # noqa: F401
11
+import vllm.model_executor.layers.mhc as mhc_layers
12
from vllm.model_executor.kernels.mhc.tilelang import (
13
_tilelang_hc_prenorm_gemm,
14
_torch_hc_prenorm_gemm,
15
)
14
-from vllm.model_executor.layers.mhc import HAS_TILELANG_MHC
16
+from vllm.model_executor.layers.mhc import (
17
+ HAS_AITER_MHC,
18
+ HAS_AITER_MHC_FUSED,
19
+ HAS_AITER_MHC_FUSED_NORM,
20
+ HAS_AITER_MHC_PRE_NORM,
21
+ HAS_TILELANG_MHC,
22
+ MHCFusedPostPreOp,
23
+ MHCPreOp,
24
+)
25
from vllm.models.deepseek_v4.nvidia.model import (
26
DeepseekV4DecoderLayer,
27
DeepseekV4Model,
@@ -291,6 +301,220 @@ def test_mhc_fused_post_pre(num_tokens, hidden_size, hc_mult):
301
torch.testing.assert_close(x, layer_input_ref, atol=1e-2, rtol=1e-2)
302
303
304
+def _rocm_mhc_inputs(num_tokens=2, hidden_size=256, hc_mult=4):
305
+ residual = torch.randn(
306
+ (num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16, device=DEVICE
307
+ )
308
+ hc_mult3 = 2 * hc_mult + hc_mult * hc_mult
309
+ fn = (
310
+ torch.randn(
311
+ (hc_mult3, hc_mult * hidden_size), dtype=torch.float32, device=DEVICE
312
+ )
313
+ * 1e-4
314
+ )
315
+ hc_scale = torch.randn((3,), dtype=torch.float32, device=DEVICE) * 0.1
316
+ hc_base = torch.randn((hc_mult3,), dtype=torch.float32, device=DEVICE) * 0.1
317
+ norm_weight = torch.randn(hidden_size, dtype=torch.bfloat16, device=DEVICE)
318
+ return residual, fn, hc_scale, hc_base, norm_weight
319
+
320
+
321
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm required")
322
+def test_mhc_pre_rocm_fallback_applies_norm(monkeypatch):
323
+ set_random_seed(0)
324
+ residual, fn, hc_scale, hc_base, norm_weight = _rocm_mhc_inputs()
325
+ rms_eps = hc_pre_eps = hc_sinkhorn_eps = norm_eps = 1e-6
326
+ sinkhorn_repeat = 20
327
+ hc_post_alpha = 1.0
328
+ ref = mhc_pre_ref(
329
+ residual,
330
+ fn,
331
+ hc_scale,
332
+ hc_base,
333
+ rms_eps,
334
+ hc_pre_eps,
335
+ hc_sinkhorn_eps,
336
+ hc_post_alpha,
337
+ sinkhorn_repeat,
338
+ )
339
+ expected_layer_input = F.rms_norm(
340
+ ref[2], (ref[2].shape[-1],), norm_weight, norm_eps
341
+ )
342
+ monkeypatch.setattr(mhc_layers, "HAS_AITER_MHC", True)
343
+ monkeypatch.setattr(mhc_layers, "HAS_AITER_MHC_PRE_NORM", False)
344
+ monkeypatch.setattr(mhc_layers, "HAS_TILELANG_MHC", False)
345
+
346
+ out = object.__new__(MHCPreOp).forward_hip(
347
+ residual,
348
+ fn,
349
+ hc_scale,
350
+ hc_base,
351
+ rms_eps,
352
+ hc_pre_eps,
353
+ hc_sinkhorn_eps,
354
+ hc_post_alpha,
355
+ sinkhorn_repeat,
356
+ norm_weight=norm_weight,
357
+ norm_eps=norm_eps,
358
+ )
359
+
360
+ torch.testing.assert_close(out[0], ref[0])
361
+ torch.testing.assert_close(out[1], ref[1])
362
+ torch.testing.assert_close(out[2], expected_layer_input)
363
+
364
+
365
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm required")
366
+def test_mhc_fused_rocm_fallback_applies_norm(monkeypatch):
367
+ set_random_seed(0)
368
+ residual, fn, hc_scale, hc_base, norm_weight = _rocm_mhc_inputs()
369
+ x = torch.randn((2, 256), dtype=torch.bfloat16, device=DEVICE)
370
+ post_layer_mix = torch.randn((2, 4, 1), dtype=torch.float32, device=DEVICE)
371
+ comb_res_mix = torch.randn((2, 4, 4), dtype=torch.float32, device=DEVICE)
372
+ rms_eps = hc_pre_eps = hc_sinkhorn_eps = norm_eps = 1e-6
373
+ sinkhorn_repeat = 20
374
+ hc_post_alpha = 1.0
375
+ residual_ref = mhc_post_ref(x, residual, post_layer_mix, comb_res_mix)
376
+ pre_ref = mhc_pre_ref(
377
+ residual_ref,
378
+ fn,
379
+ hc_scale,
380
+ hc_base,
381
+ rms_eps,
382
+ hc_pre_eps,
383
+ hc_sinkhorn_eps,
384
+ hc_post_alpha,
385
+ sinkhorn_repeat,
386
+ )
387
+ expected_layer_input = F.rms_norm(
388
+ pre_ref[2], (pre_ref[2].shape[-1],), norm_weight, norm_eps
389
+ )
390
+ monkeypatch.setattr(mhc_layers, "HAS_AITER_MHC_FUSED", False)
391
+ monkeypatch.setattr(mhc_layers, "HAS_TILELANG_MHC", False)
392
+
393
+ out = object.__new__(MHCFusedPostPreOp).forward_hip(
394
+ x,
395
+ residual,
396
+ post_layer_mix,
397
+ comb_res_mix,
398
+ fn,
399
+ hc_scale,
400
+ hc_base,
401
+ rms_eps,
402
+ hc_pre_eps,
403
+ hc_sinkhorn_eps,
404
+ hc_post_alpha,
405
+ sinkhorn_repeat,
406
+ norm_weight=norm_weight,
407
+ norm_eps=norm_eps,
408
+ )
409
+
410
+ torch.testing.assert_close(out[0], residual_ref)
411
+ torch.testing.assert_close(out[1], pre_ref[0])
412
+ torch.testing.assert_close(out[2], pre_ref[1])
413
+ torch.testing.assert_close(out[3], expected_layer_input)
414
+
415
+
416
+@pytest.mark.skipif(
417
+ not (current_platform.is_rocm() and HAS_AITER_MHC and HAS_AITER_MHC_PRE_NORM),
418
+ reason="AITER mHC with fused RMSNorm required",
419
+)
420
+def test_mhc_pre_rocm_aiter_fuses_norm():
421
+ set_random_seed(0)
422
+ residual, fn, hc_scale, hc_base, norm_weight = _rocm_mhc_inputs(
423
+ num_tokens=2, hidden_size=7168
424
+ )
425
+ rms_eps = hc_pre_eps = hc_sinkhorn_eps = norm_eps = 1e-6
426
+ sinkhorn_repeat = 20
427
+ hc_post_alpha = 1.0
428
+ ref = mhc_pre_ref(
429
+ residual,
430
+ fn,
431
+ hc_scale,
432
+ hc_base,
433
+ rms_eps,
434
+ hc_pre_eps,
435
+ hc_sinkhorn_eps,
436
+ hc_post_alpha,
437
+ sinkhorn_repeat,
438
+ )
439
+ expected_layer_input = F.rms_norm(
440
+ ref[2], (ref[2].shape[-1],), norm_weight, norm_eps
441
+ )
442
+
443
+ out = object.__new__(MHCPreOp).forward_hip(
444
+ residual,
445
+ fn,
446
+ hc_scale,
447
+ hc_base,
448
+ rms_eps,
449
+ hc_pre_eps,
450
+ hc_sinkhorn_eps,
451
+ hc_post_alpha,
452
+ sinkhorn_repeat,
453
+ norm_weight=norm_weight,
454
+ norm_eps=norm_eps,
455
+ )
456
+
457
+ torch.testing.assert_close(out[0], ref[0], atol=5e-2, rtol=1e-2)
458
+ torch.testing.assert_close(out[1], ref[1], atol=5e-2, rtol=1e-2)
459
+ torch.testing.assert_close(out[2], expected_layer_input, atol=5e-2, rtol=1e-2)
460
+
461
+
462
+@pytest.mark.skipif(
463
+ not (
464
+ current_platform.is_rocm() and HAS_AITER_MHC_FUSED and HAS_AITER_MHC_FUSED_NORM
465
+ ),
466
+ reason="AITER fused mHC with RMSNorm required",
467
+)
468
+def test_mhc_fused_rocm_aiter_fuses_norm():
469
+ set_random_seed(0)
470
+ residual, fn, hc_scale, hc_base, norm_weight = _rocm_mhc_inputs(
471
+ num_tokens=2, hidden_size=7168
472
+ )
473
+ x = torch.randn((2, 7168), dtype=torch.bfloat16, device=DEVICE)
474
+ post_layer_mix = torch.randn((2, 4, 1), dtype=torch.float32, device=DEVICE)
475
+ comb_res_mix = torch.randn((2, 4, 4), dtype=torch.float32, device=DEVICE)
476
+ rms_eps = hc_pre_eps = hc_sinkhorn_eps = norm_eps = 1e-6
477
+ sinkhorn_repeat = 20
478
+ hc_post_alpha = 1.0
479
+ residual_ref = mhc_post_ref(x, residual, post_layer_mix, comb_res_mix)
480
+ pre_ref = mhc_pre_ref(
481
+ residual_ref,
482
+ fn,
483
+ hc_scale,
484
+ hc_base,
485
+ rms_eps,
486
+ hc_pre_eps,
487
+ hc_sinkhorn_eps,
488
+ hc_post_alpha,
489
+ sinkhorn_repeat,
490
+ )
491
+ expected_layer_input = F.rms_norm(
492
+ pre_ref[2], (pre_ref[2].shape[-1],), norm_weight, norm_eps
493
+ )
494
+
495
+ out = object.__new__(MHCFusedPostPreOp).forward_hip(
496
+ x,
497
+ residual,
498
+ post_layer_mix,
499
+ comb_res_mix,
500
+ fn,
501
+ hc_scale,
502
+ hc_base,
503
+ rms_eps,
504
+ hc_pre_eps,
505
+ hc_sinkhorn_eps,
506
+ hc_post_alpha,
507
+ sinkhorn_repeat,
508
+ norm_weight=norm_weight,
509
+ norm_eps=norm_eps,
510
+ )
511
+
512
+ torch.testing.assert_close(out[0], residual_ref, atol=5e-2, rtol=1e-2)
513
+ torch.testing.assert_close(out[1], pre_ref[0], atol=5e-2, rtol=1e-2)
514
+ torch.testing.assert_close(out[2], pre_ref[1], atol=5e-2, rtol=1e-2)
515
+ torch.testing.assert_close(out[3], expected_layer_input, atol=5e-2, rtol=1e-2)
516
+
517
+
518
@pytest.mark.skipif(
519
not current_platform.is_rocm(),
520
reason="ROCm required",
tests/models/registry.py
+6
@@ -294,6 +294,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
294
"GlmMoeDsaForCausalLM": _HfExamplesInfo(
295
"zai-org/GLM-5", min_transformers_version="5.0.1", is_available_online=False
296
),
297
+ "Glm5NextForCausalLM": _HfExamplesInfo("zai-org/GLM-5.3-Flash"),
298
"GPT2LMHeadModel": _HfExamplesInfo("openai-community/gpt2"),
299
"GPTBigCodeForCausalLM": _HfExamplesInfo(
300
"bigcode/starcoder",
@@ -933,6 +934,7 @@ _MULTIMODAL_EXAMPLE_MODELS = {
934
"zai-org/GLM-OCR",
935
min_transformers_version="5.1.0",
936
),
937
+ "Glm5NextForConditionalGeneration": _HfExamplesInfo("zai-org/GLM-5.3-Flash"),
938
"H2OVLChatModel": _HfExamplesInfo(
939
"h2oai/h2ovl-mississippi-800m",
940
trust_remote_code=True,
@@ -1703,6 +1705,10 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
1705
speculative_model="zai-org/GLM-OCR",
1706
min_transformers_version="5.1.0",
1707
),
1708
+ "Glm5NextMTPModel": _HfExamplesInfo(
1709
+ "zai-org/GLM-5.3-Flash",
1710
+ speculative_model="zai-org/GLM-5.3-Flash",
1711
+ ),
1712
"HYV3MTPModel": _HfExamplesInfo(
1713
"tencent/Hy3-preview",
1714
speculative_model="tencent/Hy3-preview",
tests/models/test_initialization.py
+1
-1
@@ -142,7 +142,7 @@ def can_initialize(
142
patch.object(V1EngineCore, "_initialize_kv_caches", _initialize_kv_caches_v1),
143
monkeypatch.context() as m,
144
):
145
- if requires_spawn_multiprocessing():
145
+ if requires_spawn_multiprocessing() or model_arch == "Glm5NextForCausalLM":
146
# The EngineCore subprocess re-imports the class and does not
147
# inherit the KV-cache patch above, so it OOMs. Run in-process
148
# so the patch applies.
tests/multimodal/test_video.py
+231
@@ -20,6 +20,7 @@ from vllm.multimodal.video import (
20
PYNVVIDEOCODEC_VIDEO_BACKEND,
21
VIDEO_LOADER_REGISTRY,
22
DynamicVideoBackend,
23
+ Glm5NextVideoBackend,
24
GLM46VVideoBackend,
25
Molmo2VideoBackend,
26
Qwen2VLVideoBackend,
@@ -1470,3 +1471,233 @@ def test_glm46v_duration_estimation_from_fps():
1471
assert len(indices) > 0
1472
assert len(indices) % 2 == 0
1473
assert all(0 <= idx < 90 for idx in indices)
1474
+
1475
+
1476
+def test_glm5next_backend_selected_for_processor():
1477
+ """Glm5NextVideoProcessor maps to the glm5next loader so only the
1478
+ sampled frames are decoded instead of the whole container. Both the
1479
+ borrowed-config spelling and the dedicated Glm5next class name (landing
1480
+ with the new checkpoint) must resolve."""
1481
+ for name in ("Glm5NextVideoProcessor", "Glm5nextVideoProcessor"):
1482
+ assert VIDEO_LOADER_REGISTRY.get_backend_for_video_processor(name) == "glm5next"
1483
+
1484
+
1485
+@pytest.mark.parametrize(
1486
+ ("total_frames", "original_fps", "duration", "fps", "max_frames"),
1487
+ [
1488
+ (900, 30.0, 30.0, -1, None), # 30s at flat 2.0 raw fps -> 60 frames
1489
+ (3000, 30.0, 100.0, -1, None),
1490
+ (72000, 30.0, 2400.0, -1, None), # 2048 cap
1491
+ (48, 2.0, 24.0, -1, None),
1492
+ (7, 30.0, 10.0, -1, None), # short video -> uniform spread + dedup
1493
+ (300, 25.0, 0, -1, None), # duration derived from frame count
1494
+ (900, 30.0, 30.0, 4, None), # request fps override (raw fps)
1495
+ (900, 30.0, 30.0, -1, 16), # request max_frames override
1496
+ ],
1497
+)
1498
+def test_glm5next_backend_indices_match_sampler(
1499
+ total_frames, original_fps, duration, fps, max_frames
1500
+):
1501
+ """The loader must select exactly the frames the processor's sampler
1502
+ would, with target.fps mapping onto the raw-fps override."""
1503
+ from vllm.transformers_utils.processors.glm5next import (
1504
+ glm_sample_frame_indices,
1505
+ )
1506
+
1507
+ source = VideoSourceMetadata(
1508
+ total_frames_num=total_frames, original_fps=original_fps, duration=duration
1509
+ )
1510
+ target = VideoTargetMetadata(num_frames=-1, fps=fps, max_duration=-1)
1511
+
1512
+ indices = Glm5NextVideoBackend.compute_frames_index_to_sample(
1513
+ source, target, max_frames=max_frames
1514
+ )
1515
+
1516
+ assert indices == glm_sample_frame_indices(
1517
+ total_frames,
1518
+ original_fps,
1519
+ duration,
1520
+ target_fps=fps if fps > 0 else None,
1521
+ max_frame_count=max_frames,
1522
+ )
1523
+ assert len(indices) % 2 == 0
1524
+ assert indices == sorted(indices) # pair padding may repeat the last frame
1525
+ assert all(0 <= idx < total_frames for idx in indices)
1526
+
1527
+
1528
+def test_glm5next_backend_metadata_contract():
1529
+ """create_hf_metadata reports the subset so the processor skips
1530
+ re-sampling (do_sample_frames=False) and keeps the original totals."""
1531
+ source = VideoSourceMetadata(total_frames_num=900, original_fps=30.0, duration=30.0)
1532
+ target = VideoTargetMetadata(num_frames=-1, fps=-1, max_duration=-1)
1533
+ indices = Glm5NextVideoBackend.compute_frames_index_to_sample(source, target)
1534
+
1535
+ metadata = Glm5NextVideoBackend.create_hf_metadata(
1536
+ source, indices, video_backend="glm5next"
1537
+ )
1538
+ assert metadata["do_sample_frames"] is False
1539
+ assert metadata["frames_indices"] == indices
1540
+ assert metadata["total_num_frames"] == 900
1541
+ assert metadata["fps"] == 30.0
1542
+ assert metadata["duration"] == 30.0
1543
+
1544
+ # A fully-selected source keeps do_sample_frames=True so the processor's
1545
+ # sampler takes over on the complete frame set.
1546
+ full = list(range(48))
1547
+ assert (
1548
+ Glm5NextVideoBackend.create_hf_metadata(
1549
+ VideoSourceMetadata(total_frames_num=48, original_fps=2.0, duration=24.0),
1550
+ full,
1551
+ video_backend="glm5next",
1552
+ )["do_sample_frames"]
1553
+ is True
1554
+ )
1555
+
1556
+
1557
+def _write_gray_video(tmp_path, total_frames, fps, size=(32, 32)):
1558
+ """Synthetic clip whose frame i is flat gray level i (near-lossless under
1559
+ mp4v), so a decoded frame's level maps back to its source index."""
1560
+ cv2 = pytest.importorskip("cv2")
1561
+
1562
+ path = tmp_path / f"gray_{total_frames}_{fps}.mp4"
1563
+ writer = cv2.VideoWriter(
1564
+ str(path), cv2.VideoWriter_fourcc(*"mp4v"), fps, size, isColor=False
1565
+ )
1566
+ assert writer.isOpened()
1567
+ for i in range(total_frames):
1568
+ writer.write(np.full((*size, 1), i, dtype=np.uint8))
1569
+ writer.release()
1570
+ return path
1571
+
1572
+
1573
+class _CountingCap:
1574
+ """Proxy over a real capture that counts grab/seek decoding work."""
1575
+
1576
+ def __init__(self, cap):
1577
+ self._cap = cap
1578
+ self.grabs = 0
1579
+ self.seeks = 0
1580
+
1581
+ def get(self, prop):
1582
+ return self._cap.get(prop)
1583
+
1584
+ def grab(self):
1585
+ self.grabs += 1
1586
+ return self._cap.grab()
1587
+
1588
+ def set(self, prop, value):
1589
+ self.seeks += 1
1590
+ return self._cap.set(prop, value)
1591
+
1592
+ def read(self):
1593
+ return self._cap.read()
1594
+
1595
+
1596
+@pytest.mark.parametrize("backend", ["opencv", "torchcodec"])
1597
+def test_glm5next_backend_codec_parity(tmp_path, backend):
1598
+ """Every codec samples the same GLM indices and decodes the same
1599
+ frames; the OpenCV seek reader and torchcodec batched index-exact decode
1600
+ must agree."""
1601
+ if backend == "torchcodec":
1602
+ pytest.importorskip("torchcodec")
1603
+
1604
+ from vllm.transformers_utils.processors.glm5next import (
1605
+ glm_sample_frame_indices,
1606
+ )
1607
+
1608
+ total_frames, fps = 120, 10
1609
+ path = _write_gray_video(tmp_path, total_frames, fps)
1610
+ # Dense default sampling (gap 5) and a sparse max_frames cap (gap 20).
1611
+ for max_frames in (None, 6):
1612
+ kwargs = {} if max_frames is None else {"max_frames": max_frames}
1613
+ expected = glm_sample_frame_indices(
1614
+ total_frames, float(fps), 12.0, max_frame_count=max_frames
1615
+ )
1616
+
1617
+ frames, metadata = Glm5NextVideoBackend.load_bytes(
1618
+ path.read_bytes(), backend=backend, **kwargs
1619
+ )
1620
+
1621
+ assert metadata["frames_indices"] == expected
1622
+ assert metadata["video_backend"].startswith(backend)
1623
+ assert len(frames) == len(expected)
1624
+ for i, idx in enumerate(expected):
1625
+ assert abs(round(float(np.asarray(frames[i]).mean())) - idx) <= 1
1626
+
1627
+
1628
+def test_glm5next_backend_decodes_only_sampled_frames(tmp_path):
1629
+ """End to end over a synthetic clip: load_bytes returns exactly the
1630
+ sampler's frame count, with the right frame content at each index."""
1631
+ pytest.importorskip("cv2")
1632
+
1633
+ total_frames, fps = 60, 10
1634
+ path = _write_gray_video(tmp_path, total_frames, fps)
1635
+
1636
+ from vllm.transformers_utils.processors.glm5next import (
1637
+ glm_sample_frame_indices,
1638
+ )
1639
+
1640
+ expected = glm_sample_frame_indices(total_frames, float(fps), 6.0)
1641
+
1642
+ frames, metadata = Glm5NextVideoBackend.load_bytes(path.read_bytes())
1643
+
1644
+ assert len(frames) == len(expected)
1645
+ assert metadata["frames_indices"] == expected
1646
+ assert metadata["do_sample_frames"] is (len(expected) == total_frames)
1647
+ # Flat gray frames survive mp4v near-losslessly: each decoded frame's
1648
+ # level maps back to its source frame index.
1649
+ for frame, idx in zip(frames, expected):
1650
+ decoded_idx = round(float(np.asarray(frame).mean()))
1651
+ assert abs(decoded_idx - idx) <= 1
1652
+
1653
+
1654
+def test_glm5next_read_frames_seeks_past_large_gaps(tmp_path):
1655
+ """Sparse targets must not walk the container: the stock reader grabs
1656
+ every frame up to the last index; the GLM reader seeks instead."""
1657
+ cv2 = pytest.importorskip("cv2")
1658
+
1659
+ total_frames, fps = 200, 10
1660
+ path = _write_gray_video(tmp_path, total_frames, fps)
1661
+ targets = [0, 80, 160, 190] # gaps of 80/80/30 -> only 30 <= threshold
1662
+
1663
+ stock = cv2.VideoCapture(str(path))
1664
+ _, stock_indices = VideoBackend.read_frames(stock, targets, total_frames)
1665
+ stock.release()
1666
+ assert stock_indices == targets
1667
+
1668
+ cap = _CountingCap(cv2.VideoCapture(str(path)))
1669
+ frames, indices = Glm5NextVideoBackend.read_frames(cap, targets, total_frames)
1670
+ cap._cap.release()
1671
+
1672
+ assert indices == targets == stock_indices
1673
+ for frame, idx in zip(frames, targets):
1674
+ assert abs(round(float(np.asarray(frame).mean())) - idx) <= 1
1675
+ # Walks only the sub-threshold 30-frame hop; the two 80-frame gaps are
1676
+ # seeks. The stock reader grabs all 190 preceding frames.
1677
+ assert cap.grabs <= 29
1678
+ assert cap.seeks == 3
1679
+
1680
+
1681
+def test_glm5next_read_frames_dense_walk_matches_stock(tmp_path):
1682
+ """Dense targets keep the sequential walk (seeking would be slower) and
1683
+ return the same frames as the stock reader."""
1684
+ cv2 = pytest.importorskip("cv2")
1685
+
1686
+ total_frames, fps = 120, 10
1687
+ path = _write_gray_video(tmp_path, total_frames, fps)
1688
+ targets = list(range(0, total_frames, 15)) # gaps of 15 -> all walking
1689
+
1690
+ stock = cv2.VideoCapture(str(path))
1691
+ stock_frames, stock_indices = VideoBackend.read_frames(stock, targets, total_frames)
1692
+ stock.release()
1693
+
1694
+ cap = _CountingCap(cv2.VideoCapture(str(path)))
1695
+ frames, indices = Glm5NextVideoBackend.read_frames(cap, targets, total_frames)
1696
+ cap._cap.release()
1697
+
1698
+ assert indices == stock_indices
1699
+ for frame, stock_frame, idx in zip(frames, stock_frames, targets):
1700
+ assert abs(float(np.asarray(frame).mean()) - float(stock_frame.mean())) <= 2.0
1701
+ assert abs(round(float(np.asarray(frame).mean())) - idx) <= 1
1702
+ # One initial seek, then pure walking -- no re-seek churn.
1703
+ assert cap.seeks == 1
tests/transformers_utils/processors/test_glm5next.py
new
+389
@@ -0,0 +1,389 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Unit tests for the vLLM-native GLM-5.3-Flash multimodal processor."""
4
+
5
+import math
6
+
7
+import pytest
8
+import torch
9
+from PIL import Image
10
+from transformers.image_utils import PILImageResampling
11
+from transformers.video_utils import VideoMetadata
12
+
13
+from vllm.transformers_utils.processors.glm5next import (
14
+ Glm5NextImageProcessor,
15
+ Glm5NextVideoProcessor,
16
+ _get_pad_content_size,
17
+ _resize_or_pad,
18
+ glm_sample_frame_indices,
19
+ smart_resize,
20
+)
21
+
22
+PATCH_SIZE = 14
23
+MERGE_SIZE = 2
24
+PATCH_EXPAND_FACTOR = 1 # checkpoint processor_config.json
25
+FACTOR = PATCH_SIZE * MERGE_SIZE * PATCH_EXPAND_FACTOR # 28
26
+PIXELS_PER_TOKEN = 2 * (PATCH_SIZE * MERGE_SIZE) ** 2 # 1568
27
+MIN_TOKENS = 16
28
+MAX_TOKENS = 800 # serving-cap equivalent budget, keeps test canvases small
29
+MIN_PIXELS = MIN_TOKENS * PIXELS_PER_TOKEN
30
+MAX_PIXELS = MAX_TOKENS * PIXELS_PER_TOKEN
31
+
32
+
33
+def resize(height: int, width: int, t: int = 2) -> tuple[int, int]:
34
+ return smart_resize(
35
+ t,
36
+ height,
37
+ width,
38
+ t_factor=2,
39
+ h_factor=FACTOR,
40
+ w_factor=FACTOR,
41
+ min_pixels=MIN_PIXELS,
42
+ max_pixels=MAX_PIXELS,
43
+ )
44
+
45
+
46
+@pytest.mark.parametrize(
47
+ ("height", "width", "t", "expected"),
48
+ [
49
+ (300, 400, 2, (308, 420)), # ceil alignment, no budget pressure
50
+ (50, 80, 2, (112, 168)), # below min budget: proportional upscale
51
+ (57, 3, 2, (504, 28)), # slim side ceils to one factor
52
+ (2160, 3840, 2, (588, 1064)), # shrink: max canvas under the cap
53
+ (112, 11200, 2, (84, 7420)), # extreme aspect stays proportional
54
+ (300, 400, 16, (252, 308)), # frame count eats into the budget
55
+ (112, 11200, 16, (28, 2800)), # slim side stays positive, budget holds
56
+ ],
57
+)
58
+def test_smart_resize_reference_values(height, width, t, expected):
59
+ got = smart_resize(
60
+ t,
61
+ height,
62
+ width,
63
+ t_factor=2,
64
+ h_factor=FACTOR,
65
+ w_factor=FACTOR,
66
+ min_pixels=MIN_PIXELS,
67
+ max_pixels=MAX_PIXELS,
68
+ )
69
+ assert got == expected
70
+ # The aligned canvas never exceeds the budget it was fitted against.
71
+ t_bar = max(2, round(t / 2) * 2)
72
+ assert t_bar * got[0] * got[1] <= MAX_PIXELS
73
+
74
+
75
+def test_smart_resize_stays_snapped_and_positive():
76
+ heights = (1, 27, 57, 111, 113, 300, 720, 2160)
77
+ widths = (1, 27, 57, 111, 113, 400, 1280, 11200)
78
+ for h in heights:
79
+ for w in widths:
80
+ for t in (2, 4, 16, 64):
81
+ resized_h, resized_w = resize(h, w, t)
82
+ assert resized_h > 0 and resized_w > 0
83
+ assert resized_h % FACTOR == 0
84
+ assert resized_w % FACTOR == 0
85
+ t_bar = max(2, round(t / 2) * 2)
86
+ assert t_bar * resized_h * resized_w <= MAX_PIXELS
87
+
88
+
89
+@pytest.mark.parametrize(
90
+ ("height", "width", "match"),
91
+ [
92
+ (0, 100, "must be positive"),
93
+ (100, 0, "must be positive"),
94
+ ],
95
+)
96
+def test_smart_resize_rejects_degenerate_inputs(height, width, match):
97
+ with pytest.raises(ValueError, match=match):
98
+ resize(height, width, 2)
99
+
100
+
101
+def test_smart_resize_rejects_inverted_budget():
102
+ with pytest.raises(ValueError, match="min_pixels must be less than or equal"):
103
+ smart_resize(
104
+ 2,
105
+ 100,
106
+ 100,
107
+ t_factor=2,
108
+ h_factor=FACTOR,
109
+ w_factor=FACTOR,
110
+ min_pixels=MAX_PIXELS,
111
+ max_pixels=MIN_PIXELS,
112
+ )
113
+
114
+
115
+def test_get_pad_content_size():
116
+ # Oversized content shrinks proportionally, never upscales by default.
117
+ assert _get_pad_content_size(300, 400, 308, 420) == (300, 400)
118
+ assert _get_pad_content_size(600, 800, 308, 420) == (308, 410)
119
+ # allow_upscale enlarges small content toward the canvas.
120
+ assert _get_pad_content_size(28, 28, 112, 112, allow_upscale=True) == (112, 112)
121
+
122
+
123
+def test_resize_or_pad_pads_right_and_bottom():
124
+ stacked = torch.rand(1, 3, 300, 400)
125
+
126
+ def identity_resize(x, size, resample=None):
127
+ assert (size.height, size.width) == (300, 400)
128
+ return x
129
+
130
+ padded = _resize_or_pad(stacked, 308, 420, "pad", None, identity_resize)
131
+ assert padded.shape == (1, 3, 308, 420)
132
+ torch.testing.assert_close(padded[..., :300, :400], stacked)
133
+ assert padded[..., 300:, :].eq(0).all()
134
+ assert padded[..., :, 400:].eq(0).all()
135
+
136
+ def force_resize(x, size, resample=None):
137
+ return torch.zeros(x.shape[0], x.shape[1], size.height, size.width)
138
+
139
+ assert _resize_or_pad(stacked, 308, 420, "resize", None, force_resize).shape == (
140
+ 1,
141
+ 3,
142
+ 308,
143
+ 420,
144
+ )
145
+ with pytest.raises(ValueError, match="resize_mode"):
146
+ _resize_or_pad(stacked, 308, 420, "crop", None, identity_resize)
147
+
148
+
149
+@pytest.fixture(scope="module")
150
+def image_processor():
151
+ return Glm5NextImageProcessor(
152
+ min_image_tokens=MIN_TOKENS, max_image_tokens=MAX_TOKENS
153
+ )
154
+
155
+
156
+@pytest.mark.parametrize(
157
+ ("height", "width"),
158
+ [(50, 80), (300, 400), (2160, 3840), (112, 11200)],
159
+)
160
+def test_image_grid_matches_smart_resize(image_processor, height, width):
161
+ out = image_processor(Image.new("RGB", (width, height)), return_tensors="pt")
162
+ resized_h, resized_w = resize(height, width)
163
+ grid = out["image_grid_thw"][0].tolist()
164
+ assert grid == [1, resized_h // PATCH_SIZE, resized_w // PATCH_SIZE]
165
+ assert out["pixel_values"].shape == (
166
+ grid[1] * grid[2],
167
+ 3 * 2 * PATCH_SIZE * PATCH_SIZE,
168
+ )
169
+
170
+
171
+@pytest.mark.parametrize(
172
+ ("height", "width"),
173
+ [(50, 80), (300, 400), (2160, 3840), (112, 11200)],
174
+)
175
+def test_image_patch_count_matches_preprocess(image_processor, height, width):
176
+ out = image_processor(Image.new("RGB", (width, height)), return_tensors="pt")
177
+ grid = out["image_grid_thw"][0].tolist()
178
+ expected = grid[1] * grid[2]
179
+ assert image_processor.get_number_of_image_patches(height, width) == expected
180
+
181
+
182
+@pytest.fixture(scope="module")
183
+def video_processor():
184
+ return Glm5NextVideoProcessor(
185
+ min_image_tokens=MIN_TOKENS, max_image_tokens=MAX_TOKENS
186
+ )
187
+
188
+
189
+def run_video_preprocess(video_processor, frames, **overrides):
190
+ kwargs = dict(
191
+ resample=PILImageResampling.BICUBIC,
192
+ image_mean=(0.48145466, 0.4578275, 0.40821073),
193
+ image_std=(0.26862954, 0.26130258, 0.27577711),
194
+ patch_size=PATCH_SIZE,
195
+ temporal_patch_size=2,
196
+ merge_size=MERGE_SIZE,
197
+ patch_expand_factor=PATCH_EXPAND_FACTOR,
198
+ return_tensors="pt",
199
+ )
200
+ kwargs.update(overrides)
201
+ return video_processor._preprocess([frames], **kwargs)
202
+
203
+
204
+def test_video_preprocess_pads_odd_frame_count(video_processor):
205
+ out = run_video_preprocess(video_processor, torch.rand(7, 3, 300, 400))
206
+ grid = out["video_grid_thw"][0].tolist()
207
+ # 7 frames padded to 8 -> grid_t 4; canvas ceil-aligned (308, 420).
208
+ assert grid == [4, 22, 30]
209
+ assert out["pixel_values_videos"].shape == (
210
+ 4 * 22 * 30,
211
+ 3 * 2 * PATCH_SIZE * PATCH_SIZE,
212
+ )
213
+
214
+
215
+def test_video_preprocess_keeps_min_side_under_frame_budget(video_processor):
216
+ out = run_video_preprocess(video_processor, torch.rand(16, 3, 112, 11200))
217
+ # Budget-fitted canvas (28, 2800): 16 * 28 * 2800 == the pixel cap.
218
+ assert out["video_grid_thw"][0].tolist() == [8, 2, 200] # height 28, not 0
219
+
220
+
221
+def test_video_preprocess_pads_content_not_distort(video_processor):
222
+ frames = torch.zeros(4, 3, 300, 400)
223
+ frames[..., 50, 60] = 1.0 # marker inside the content area
224
+ # Bypass rescale/normalize so zero padding stays exactly zero.
225
+ out = run_video_preprocess(
226
+ video_processor, frames, do_rescale=False, do_normalize=False
227
+ )
228
+ # Pad mode keeps the 300x400 content aspect on the (308, 420) canvas
229
+ # with zero padding on the right/bottom.
230
+ grid = out["video_grid_thw"][0].tolist()
231
+ assert grid == [2, 22, 30]
232
+ patches = out["pixel_values_videos"].view(
233
+ grid[0],
234
+ grid[1] // MERGE_SIZE,
235
+ grid[2] // MERGE_SIZE,
236
+ MERGE_SIZE,
237
+ MERGE_SIZE,
238
+ 3 * 2 * PATCH_SIZE * PATCH_SIZE,
239
+ )
240
+ frame = patches[0].permute(0, 3, 1, 4, 2).reshape(grid[1], grid[2], -1)
241
+ # Patch columns fully right of the 400px content are pure padding.
242
+ assert frame[:, math.ceil(400 / PATCH_SIZE) :].abs().max() == 0
243
+
244
+
245
+@pytest.mark.parametrize(
246
+ ("total_frames", "fps", "duration", "expected_len", "first", "last"),
247
+ [
248
+ # Dense source: the tp-scaled greedy overshoots extract_t, so the
249
+ # fixup spreads picks uniformly across the whole video (linspace).
250
+ (900, 30.0, 30.0, 60, 0, 899),
251
+ (3000, 30.0, 100.0, 200, 0, 2999),
252
+ # extract_t capped at 2048 frames.
253
+ (72000, 30.0, 2400.0, 2048, 0, 71999),
254
+ # Low container fps: the greedy picks every frame.
255
+ (48, 2.0, 24.0, 48, 0, 47),
256
+ # Duration derived from the frame count when metadata lacks it.
257
+ (300, 25.0, 0, 26, 0, 299),
258
+ ],
259
+)
260
+def test_glm_sample_frame_indices_behaviour(
261
+ total_frames, fps, duration, expected_len, first, last
262
+):
263
+ indices = glm_sample_frame_indices(total_frames, fps, duration)
264
+ assert len(indices) == expected_len
265
+ assert indices[0] == first
266
+ assert indices[-1] == last
267
+ assert len(indices) % 2 == 0
268
+ assert indices == sorted(indices)
269
+ assert all(0 <= idx < total_frames for idx in indices)
270
+
271
+
272
+def test_glm_sample_frame_indices_short_clip_floor_spread():
273
+ # extract_t (20) > total (7): evenly spaced timestamps, deduplicated,
274
+ # then pair-padded to an even count.
275
+ assert glm_sample_frame_indices(7, 30.0, 10.0) == [0, 1, 2, 3, 4, 5, 6, 6]
276
+
277
+
278
+@pytest.mark.parametrize(
279
+ ("kwargs", "expected_len"),
280
+ [
281
+ ({"target_fps": 0.5}, 16),
282
+ ({"max_frame_count": 16}, 16),
283
+ ({"target_fps": 8}, 240), # 30s * 8 = 240, under the 2048 cap
284
+ ],
285
+)
286
+def test_glm_sample_frame_indices_request_overrides(kwargs, expected_len):
287
+ indices = glm_sample_frame_indices(900, 30.0, 30.0, **kwargs)
288
+ assert len(indices) == expected_len
289
+ assert len(indices) % 2 == 0
290
+
291
+
292
+def test_sample_frames_fps_interval(video_processor):
293
+ indices = video_processor.sample_frames(
294
+ VideoMetadata(total_num_frames=900, fps=30.0, duration=30.0)
295
+ )
296
+ assert len(indices) == 60
297
+ assert indices[0] == 0
298
+ assert indices[-1] == 899 # uniform spread reaches the final frame
299
+
300
+ # Request overrides reach the sampler through both kwarg spellings. The
301
+ # 15-frame spread is odd, so pair-padding duplicates the last frame.
302
+ assert (
303
+ len(
304
+ video_processor.sample_frames(
305
+ VideoMetadata(total_num_frames=900, fps=30.0, duration=30.0),
306
+ fps=0.5,
307
+ )
308
+ )
309
+ == 16
310
+ )
311
+ assert (
312
+ len(
313
+ video_processor.sample_frames(
314
+ VideoMetadata(total_num_frames=900, fps=30.0, duration=30.0),
315
+ target_fps=0.5,
316
+ )
317
+ )
318
+ == 16
319
+ )
320
+ assert (
321
+ len(
322
+ video_processor.sample_frames(
323
+ VideoMetadata(total_num_frames=900, fps=30.0, duration=30.0),
324
+ max_frames=16,
325
+ )
326
+ )
327
+ == 16
328
+ )
329
+
330
+
331
+def test_defaults_mirror_checkpoint_config():
332
+ """Bare instantiation matches the checkpoint's ``processor_config.json``
333
+ token-budget style defaults; ``size`` carries no budget."""
334
+ image_processor = Glm5NextImageProcessor()
335
+ assert image_processor.patch_expand_factor == 1
336
+ assert image_processor.min_image_tokens == 16
337
+ assert image_processor.max_image_tokens == 8000
338
+ assert image_processor.resize_mode == "pad"
339
+
340
+ video_processor = Glm5NextVideoProcessor()
341
+ assert video_processor.patch_expand_factor == 1
342
+ assert video_processor.min_image_tokens == 16
343
+ assert video_processor.max_image_tokens == 240000
344
+ assert video_processor.fps_interval == 2.0
345
+ assert video_processor.max_frame_count_dynamic == 2048
346
+
347
+
348
+def test_token_budgets_drive_geometry():
349
+ """The token bounds fully determine the pixel budget and the grid."""
350
+ token_proc = Glm5NextImageProcessor(min_image_tokens=64, max_image_tokens=512)
351
+ img = Image.new("RGB", (400, 300))
352
+ grid = token_proc(img, return_tensors="pt")["image_grid_thw"][0].tolist()
353
+ # Budgets 100,352..802,816 px: the 300x400 canvas ceils to (308, 420)
354
+ # without hitting either bound -> 22x30 patches.
355
+ assert grid == [1, 22, 30]
356
+ assert token_proc.get_number_of_image_patches(300, 400) == 22 * 30
357
+
358
+
359
+def test_missing_token_budgets_rejected():
360
+ proc = Glm5NextImageProcessor(min_image_tokens=None)
361
+ with pytest.raises(ValueError, match="min_image_tokens"):
362
+ proc(Image.new("RGB", (400, 300)), return_tensors="pt")
363
+
364
+
365
+def test_video_config_fields_land():
366
+ """fps_interval / max_frame_count_dynamic from the dedicated config
367
+ shape sampling without any request overrides."""
368
+ proc = Glm5NextVideoProcessor(
369
+ fps_interval=4,
370
+ max_frame_count_dynamic=32,
371
+ )
372
+ indices = proc.sample_frames(
373
+ VideoMetadata(total_num_frames=900, fps=30.0, duration=30.0)
374
+ )
375
+ # 30s * 4 = 120 candidates, capped at 32 -> uniform spread of 32.
376
+ assert len(indices) == 32
377
+ assert indices[0] == 0
378
+ assert indices[-1] == 899
379
+
380
+ # Request overrides still win over the config values.
381
+ assert (
382
+ len(
383
+ proc.sample_frames(
384
+ VideoMetadata(total_num_frames=900, fps=30.0, duration=30.0),
385
+ fps=0.5,
386
+ )
387
+ )
388
+ == 16
389
+ )
tests/transformers_utils/test_config.py
+54
@@ -10,6 +10,7 @@ from types import SimpleNamespace
10
from typing import cast
11
from unittest.mock import MagicMock, patch
12
13
+import pytest
14
from transformers import PretrainedConfig
15
16
from vllm.config.model import ModelConfig
@@ -19,6 +20,59 @@ from vllm.transformers_utils.config import (
20
get_safetensors_params_metadata,
21
try_get_generation_config,
22
)
23
+from vllm.transformers_utils.configs.glm5_next import (
24
+ Glm5NextConfig,
25
+ Glm5NextTextConfig,
26
+ Glm5NextVisionConfig,
27
+)
28
+
29
+
30
+def test_glm5_next_accepts_deepseek_sparse_attention_layers():
31
+ layer_types = ["linear_attention", "deepseek_sparse_attention"]
32
+
33
+ config = Glm5NextTextConfig(
34
+ num_hidden_layers=len(layer_types), layer_types=layer_types
35
+ )
36
+
37
+ assert config.layer_types == layer_types
38
+ assert config.layers_block_type == ["linear_attention", "attention"]
39
+
40
+
41
+def test_glm5_next_accepts_prebuilt_subconfigs():
42
+ text_config = Glm5NextTextConfig(hidden_size=1024)
43
+ vision_config = Glm5NextVisionConfig(hidden_size=768)
44
+
45
+ config = Glm5NextConfig(
46
+ text_config=text_config,
47
+ vision_config=vision_config,
48
+ )
49
+
50
+ assert config.text_config is text_config
51
+ assert config.vision_config is vision_config
52
+
53
+
54
+@pytest.mark.parametrize(
55
+ ("kwargs", "option"),
56
+ [
57
+ (
58
+ {"index_topk": 2048, "index_dsa_use_layernorm": False},
59
+ "index_dsa_use_layernorm",
60
+ ),
61
+ (
62
+ {"index_topk": 2048, "index_kpool_compress": False},
63
+ "index_kpool_compress",
64
+ ),
65
+ (
66
+ {"index_topk": 2048, "index_kpool_always_select_tail": False},
67
+ "index_kpool_always_select_tail",
68
+ ),
69
+ ({"hres_vwnstyle": False}, "hres_vwnstyle"),
70
+ ({"mhc_no_norm_weight": True}, "mhc_no_norm_weight"),
71
+ ],
72
+)
73
+def test_glm5_next_rejects_unimplemented_config_options(kwargs, option):
74
+ with pytest.raises(NotImplementedError, match=option):
75
+ Glm5NextTextConfig(**kwargs)
76
77
78
def test_get_llama3_eos_token():
tests/utils.py
+7
-6
@@ -22,7 +22,7 @@ from collections.abc import Callable, Iterable, MutableMapping, Sequence
22
from contextlib import ExitStack, contextmanager
23
from multiprocessing import Process, get_context
24
from pathlib import Path
25
-from typing import Any, Literal, cast
25
+from typing import TYPE_CHECKING, Any, Literal, cast
26
from unittest.mock import patch
27
28
import anthropic
@@ -46,10 +46,6 @@ from vllm.distributed import (
46
from vllm.engine.arg_utils import AsyncEngineArgs
47
from vllm.entrypoints.cli.serve import ServeSubcommand
48
from vllm.logger import init_logger
49
-from vllm.model_executor.kernels.linear import (
50
- _KernelT,
51
- init_fp8_linear_kernel,
52
-)
49
from vllm.model_executor.layers.quantization.utils.quant_utils import (
50
QuantKey,
51
)
@@ -65,6 +61,9 @@ from vllm.utils.torch_utils import (
61
)
62
from vllm.v1.engine.utils import get_engine_process_shutdown_timeout
63
64
+if TYPE_CHECKING:
65
+ from vllm.model_executor.kernels.linear import _KernelT
66
+
67
logger = init_logger(__name__)
68
69
FP8_DTYPE = current_platform.fp8_dtype()
@@ -2349,9 +2348,11 @@ class TestFP8Layer(torch.nn.Module):
2348
out_dtype: torch.dtype | None = None,
2349
transpose_weights: bool = False,
2350
device: torch.device | None = None,
2352
- force_kernel: type[_KernelT] | None = None,
2351
+ force_kernel: "type[_KernelT] | None" = None,
2352
):
2353
super().__init__()
2354
+ from vllm.model_executor.kernels.linear import init_fp8_linear_kernel
2355
+
2356
self.input_size_per_partition = weight_shape[1]
2357
self.output_size_per_partition = weight_shape[0]
2358
self.logical_widths = [self.output_size_per_partition]
tests/v1/attention/test_cuda_backend_probe_errors.py
+43
-1
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Tests for error handling in the CUDA attention backend probe.
3
+"""Tests for CUDA attention backend selection and probe error handling.
4
5
Environment-shaped probe failures (missing packages, unreadable caches,
6
broken driver installs) must mark the backend unavailable rather than
@@ -89,6 +89,48 @@ def test_get_valid_backends_keeps_probing_after_failure():
89
assert valid
90
91
92
+def test_sm90_nope_mla_prefers_flashinfer_without_changing_rope_order():
93
+ backend_cls = MagicMock()
94
+ backend_cls.validate_configuration.return_value = []
95
+ sparse_backends = {
96
+ AttentionBackendEnum.FLASH_ATTN_MLA_SPARSE,
97
+ AttentionBackendEnum.FLASHMLA_SPARSE,
98
+ AttentionBackendEnum.FLASHINFER_MLA_SPARSE_SM90,
99
+ }
100
+
101
+ def sparse_order(head_size: int) -> list[AttentionBackendEnum]:
102
+ config = SELECTOR_CONFIG._replace(
103
+ head_size=head_size,
104
+ use_mla=True,
105
+ use_sparse=True,
106
+ )
107
+ with patch(
108
+ "vllm.platforms.cuda._get_attn_backend_class",
109
+ return_value=backend_cls,
110
+ ):
111
+ valid, _ = CudaPlatform.get_valid_backends(
112
+ device_capability=SM90,
113
+ attn_selector_config=config,
114
+ num_heads=32,
115
+ )
116
+ return [
117
+ candidate.backend
118
+ for candidate in valid
119
+ if candidate.backend in sparse_backends
120
+ ]
121
+
122
+ assert sparse_order(512) == [
123
+ AttentionBackendEnum.FLASHINFER_MLA_SPARSE_SM90,
124
+ AttentionBackendEnum.FLASH_ATTN_MLA_SPARSE,
125
+ AttentionBackendEnum.FLASHMLA_SPARSE,
126
+ ]
127
+ assert sparse_order(576) == [
128
+ AttentionBackendEnum.FLASH_ATTN_MLA_SPARSE,
129
+ AttentionBackendEnum.FLASHMLA_SPARSE,
130
+ AttentionBackendEnum.FLASHINFER_MLA_SPARSE_SM90,
131
+ ]
132
+
133
+
134
def test_selected_backend_probe_failure_raises_value_error_with_cause():
135
exc = OSError("libcuda.so.1: cannot open shared object file")
136
with (
tests/v1/attention/test_flashinfer_mla_sparse_sm90.py
new
+280
@@ -0,0 +1,280 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""CPU tests for the FlashInfer SM90 sparse MLA backend wiring (no GPU).
4
+
5
+The FlashInfer wrapper and top-k conversion are replaced by CPU recorders;
6
+the tests pin the contract between the impl and the kernel API: page_size=1
7
+varlen rows, reserved-buffer refresh, plan parameters (dims, NoPE/rope scale,
8
+causality), ckv/kpe cache splitting, and the backend's model-shape gates.
9
+"""
10
+
11
+from types import SimpleNamespace
12
+
13
+import pytest
14
+import torch
15
+
16
+# isort: off
17
+import vllm.v1.attention.backends.mla.flashinfer_mla_sparse_sm90 as sm90_mod
18
+from vllm.v1.attention.backends.mla.flashinfer_mla_sparse_sm90 import (
19
+ FlashInferMLASparseSM90Backend,
20
+ FlashInferMLASparseSM90Builder,
21
+ FlashInferMLASparseSM90Impl,
22
+)
23
+# isort: on
24
+
25
+BLOCK_SIZE = 64
26
+HEAD = 512
27
+TOPK = 128 # triton convert requires width % 128 == 0
28
+
29
+
30
+def ref_convert(req_id, block_table, token_indices, BLOCK_SIZE=64, **_):
31
+ out = torch.full_like(token_indices, -1)
32
+ counts = torch.zeros(token_indices.shape[0], dtype=torch.int32)
33
+ for t in range(token_indices.shape[0]):
34
+ vals = []
35
+ for j in range(token_indices.shape[1]):
36
+ pos = int(token_indices[t, j])
37
+ if pos == -1:
38
+ continue
39
+ blk = int(block_table[int(req_id[t]), pos // BLOCK_SIZE])
40
+ if blk < 0:
41
+ continue
42
+ vals.append(blk * BLOCK_SIZE + pos % BLOCK_SIZE)
43
+ out[t, : len(vals)] = torch.tensor(vals, dtype=out.dtype)
44
+ counts[t] = len(vals)
45
+ return out, counts
46
+
47
+
48
+class FakeWrapper:
49
+ def __init__(self):
50
+ self.plan_args = None
51
+ self.run_args = None
52
+
53
+ def plan(self, *args, **kwargs):
54
+ self.plan_args = (args, kwargs)
55
+
56
+ def run(self, *args, **kwargs):
57
+ q_nope, q_pe, ckv, kpe = args
58
+ self.run_args = (q_nope, q_pe, ckv, kpe, kwargs)
59
+ return torch.zeros(
60
+ q_nope.shape[0], q_nope.shape[1], ckv.shape[-1], dtype=torch.bfloat16
61
+ )
62
+
63
+
64
+class FakeState:
65
+ def __init__(self, width, max_tokens=64):
66
+ self.kv_indices = torch.zeros(max_tokens * width, dtype=torch.int32)
67
+ self.kv_len_arr = torch.zeros(max_tokens, dtype=torch.int32)
68
+ self.wrapper = FakeWrapper()
69
+ self.plan_calls = []
70
+
71
+ def plan(self, num_tokens, kv_lens):
72
+ self.plan_calls.append((num_tokens, kv_lens))
73
+
74
+
75
+def make_impl(qk_rope, kv_dtype="fp8_e4m3", num_heads=2, topk_width=TOPK):
76
+ impl = object.__new__(FlashInferMLASparseSM90Impl)
77
+ impl.num_heads = num_heads
78
+ impl.head_size = HEAD + qk_rope
79
+ impl.scale = (HEAD + qk_rope) ** -0.5
80
+ impl.kv_lora_rank = HEAD
81
+ impl.qk_rope_head_dim = qk_rope
82
+ impl.kv_cache_dtype = kv_dtype
83
+ impl.use_fp8_kv_cache = kv_dtype in ("fp8", "fp8_e4m3")
84
+ rows = 4
85
+ impl.topk_indices_buffer = torch.full((rows, topk_width), -1, dtype=torch.int32)
86
+ return impl, rows
87
+
88
+
89
+def make_batch(rows, topk_rows, own_blocks):
90
+ req_id = torch.tensor([0] * rows, dtype=torch.int32)
91
+ block_table = torch.zeros(1, 16, dtype=torch.int32)
92
+ block_table[:, 0] = own_blocks[0]
93
+ topk = torch.full((rows, TOPK), -1, dtype=torch.int32)
94
+ for t, row in enumerate(topk_rows):
95
+ topk[t, : len(row)] = torch.tensor(row, dtype=torch.int32)
96
+ return SimpleNamespace(
97
+ req_id_per_token=req_id, block_table=block_table, block_size=BLOCK_SIZE
98
+ )
99
+
100
+
101
+@pytest.mark.parametrize("qk_rope,kv_dtype", [(0, "fp8_e4m3"), (64, "auto")])
102
+def test_forward_wiring(monkeypatch, qk_rope, kv_dtype):
103
+ impl, rows = make_impl(qk_rope, kv_dtype)
104
+ state = FakeState(TOPK)
105
+ monkeypatch.setattr(
106
+ sm90_mod, "triton_convert_req_index_to_global_index", ref_convert
107
+ )
108
+
109
+ # req with context 10 < topk: 8 valid + -1 padding.
110
+ topk_rows = [
111
+ [7, 3, 1, 9, 0, 2, 5, 8] + [-1] * (TOPK - 8),
112
+ [4, 0, 2, 3, 1] + [-1] * (TOPK - 5),
113
+ [6, 5, 4] + [-1] * (TOPK - 3),
114
+ [2] + [-1] * (TOPK - 1),
115
+ ]
116
+ meta = make_batch(rows, topk_rows, [3])
117
+ meta.state = state
118
+ q_nope = torch.randn(rows, impl.num_heads, HEAD)
119
+ q_rope = torch.randn(rows, impl.num_heads, qk_rope)
120
+ cache = torch.zeros(
121
+ 8 * BLOCK_SIZE,
122
+ impl.head_size,
123
+ dtype=torch.uint8 if impl.use_fp8_kv_cache else torch.bfloat16,
124
+ )
125
+
126
+ out, lse = impl.forward_mqa(
127
+ (q_nope, q_rope), cache, meta, SimpleNamespace(_k_scale_float=0.5)
128
+ )
129
+ assert lse is None and out.shape == (rows, impl.num_heads, HEAD)
130
+
131
+ # Reserved buffers carry this step's slots; lengths are NOT refreshed
132
+ # here (the builder plans them host-side before capture/replay).
133
+ ref_slots, ref_counts = ref_convert(
134
+ meta.req_id_per_token, meta.block_table, impl.topk_indices_buffer
135
+ )
136
+ width = TOPK
137
+ got_slots = state.kv_indices[: rows * width].view(rows, width)
138
+ for t in range(rows):
139
+ k = int(ref_counts[t])
140
+ assert got_slots[t, :k].tolist() == ref_slots[t, :k].tolist()
141
+ assert state.plan_calls == []
142
+
143
+ assert state.wrapper.run_args is not None
144
+ q_pe, ckv, kpe, kwargs = state.wrapper.run_args[1:]
145
+ assert q_pe.shape == (rows, impl.num_heads, qk_rope)
146
+ assert ckv.shape == (8 * BLOCK_SIZE, 1, HEAD)
147
+ assert kpe.shape[-1] == qk_rope
148
+ if impl.use_fp8_kv_cache:
149
+ assert kwargs["ckv_scale"] == 0.5 and kwargs["kpe_scale"] == 1.0
150
+ else:
151
+ assert kwargs == {}
152
+
153
+
154
+def test_builder_attaches_its_state(monkeypatch):
155
+ builder = object.__new__(FlashInferMLASparseSM90Builder)
156
+ builder._index_topk = 2048
157
+ builder._index_kpool = 4
158
+ builder._async_scheduling = False
159
+ builder.state = FakeState(TOPK)
160
+ metadata = object.__new__(sm90_mod.FlashInferMLASparseSM90Metadata)
161
+ metadata.state = None
162
+ monkeypatch.setattr(
163
+ sm90_mod.FlashInferMLASparseMetadataBuilder,
164
+ "build",
165
+ lambda *_args, **_kwargs: metadata,
166
+ )
167
+ cam = SimpleNamespace(
168
+ num_reqs=1,
169
+ query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.int32),
170
+ seq_lens=torch.tensor([1], dtype=torch.int32),
171
+ seq_lens_cpu_upper_bound=torch.tensor([1], dtype=torch.int32),
172
+ positions=None,
173
+ )
174
+
175
+ result = builder.build(0, cam)
176
+
177
+ assert result.state is builder.state
178
+ assert builder.state.plan_calls[0][0] == 1
179
+ assert builder.state.plan_calls[0][1].tolist() == [1]
180
+
181
+
182
+def test_plan_uses_state_params(monkeypatch):
183
+ """The NoPE/rope dims and scale live on the builder state, not the layer.
184
+
185
+ plan() takes exact per-row KV lengths; the schedule is rebuilt on every
186
+ call (contexts grow between steps) and the indptrs are always full-size
187
+ with zero-query padding rows past num_tokens.
188
+ """
189
+ impl, rows = make_impl(64, "auto")
190
+ wrapper = FakeWrapper()
191
+ state = sm90_mod._SM90State.__new__(sm90_mod._SM90State)
192
+ state.device = torch.device("cpu")
193
+ state.wrapper = wrapper
194
+ state.num_heads = 4
195
+ state.kv_dtype = torch.bfloat16
196
+ state.kv_lora_rank = HEAD
197
+ state.qk_rope_head_dim = 64
198
+ state.sm_scale = 576**-0.5
199
+ state.max_tokens = 4
200
+ state.topk_width = TOPK
201
+ state.kv_indices = torch.zeros(4 * TOPK)
202
+ state._arange_cpu = torch.arange(5, dtype=torch.int32)
203
+ state._qo_cpu = torch.empty(5, dtype=torch.int32)
204
+ state._kv_cpu = torch.empty(5, dtype=torch.int32)
205
+ state._lens_cpu = torch.full((4,), TOPK, dtype=torch.int32)
206
+
207
+ state.plan(3, torch.tensor([2, 5, 7], dtype=torch.int32))
208
+ assert wrapper.plan_args is not None
209
+ args, kwargs = wrapper.plan_args
210
+ (qo, kv, indices, kv_len, heads, ckv, kpe, page, causal, scale) = args
211
+ assert qo.tolist() == [0, 1, 2, 3, 3] # clamp: rows past 3 have no queries
212
+ assert kv.tolist() == [i * TOPK for i in (0, 1, 2, 3, 3)]
213
+ assert kv_len.tolist() == [2, 5, 7, TOPK] # padded row keeps full width
214
+ assert (heads, ckv, kpe, page, causal) == (4, HEAD, 64, 1, False)
215
+ assert scale == 576**-0.5
216
+ assert kwargs["q_data_type"] == torch.bfloat16
217
+ assert kwargs["kv_data_type"] == torch.bfloat16
218
+
219
+
220
+def test_kv_lens_host_formula():
221
+ """Per-row host lengths: context == position + 1; capped at
222
+ index_topk + trailing-pool remainder past the sparse threshold."""
223
+ builder = object.__new__(FlashInferMLASparseSM90Builder)
224
+ builder._index_topk = 2048
225
+ builder._index_kpool = 4
226
+ builder._async_scheduling = False
227
+ cam = SimpleNamespace(
228
+ num_reqs=3,
229
+ query_start_loc_cpu=torch.tensor([0, 5, 7, 10], dtype=torch.int32),
230
+ seq_lens=torch.tensor([100, 9, 3000], dtype=torch.int32),
231
+ seq_lens_cpu_upper_bound=torch.tensor([100, 9, 3000], dtype=torch.int32),
232
+ positions=None,
233
+ )
234
+ num_rows, lens = builder._kv_lens_host(cam)
235
+ assert num_rows == 10
236
+ # req0: positions 95..99 -> ctx 96..100 (all <= 2048: full context)
237
+ # req1: positions 7,8 -> ctx 8,9
238
+ # req2: positions 2997..2999 -> ctx 2998..3000 (> 2048: topk + ctx%4)
239
+ assert lens.tolist() == [96, 97, 98, 99, 100, 8, 9, 2050, 2051, 2048]
240
+
241
+
242
+def test_kv_lens_host_empty():
243
+ builder = object.__new__(FlashInferMLASparseSM90Builder)
244
+ builder._index_topk = 2048
245
+ builder._index_kpool = 4
246
+ cam = SimpleNamespace(
247
+ num_reqs=0,
248
+ query_start_loc_cpu=torch.tensor([0], dtype=torch.int32),
249
+ seq_lens=torch.zeros(0, dtype=torch.int32),
250
+ )
251
+ num_rows, lens = builder._kv_lens_host(cam)
252
+ assert num_rows == 0 and lens.numel() == 0
253
+
254
+
255
+def test_supports_combination_gates(monkeypatch, default_vllm_config):
256
+ monkeypatch.setattr(sm90_mod, "has_flashinfer_sm90_nope_mla", lambda: True)
257
+ call = lambda **kw: FlashInferMLASparseSM90Backend.supports_combination(
258
+ head_size=576,
259
+ dtype=torch.bfloat16,
260
+ kv_cache_dtype="fp8_e4m3",
261
+ block_size=64,
262
+ use_mla=True,
263
+ has_sink=False,
264
+ use_sparse=True,
265
+ use_mm_prefix=False,
266
+ device_capability=SimpleNamespace(major=9),
267
+ **kw,
268
+ )
269
+ assert call() is None # no model config: only the feature gate applies
270
+
271
+ import vllm.config as cfg
272
+
273
+ monkeypatch.setattr(
274
+ cfg,
275
+ "get_current_vllm_config",
276
+ lambda: SimpleNamespace(model_config=None),
277
+ )
278
+ assert call() is None
279
+ monkeypatch.setattr(sm90_mod, "has_flashinfer_sm90_nope_mla", lambda: False)
280
+ assert "requires FlashInfer" in (call() or "")
tests/v1/attention/test_kpool_tail_slot_mapping.py
new
+369
@@ -0,0 +1,369 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""CPU tests for the kpool tail slot mapping (no GPU required).
4
+
5
+The kpool tail cache is a 1-block-per-request circular ring addressed by
6
+``pos % kpool`` (``KpoolTailSpec`` / ``KpoolTailManager``: exactly one block
7
+allocated per request, never grown, so only column 0 of its block table is
8
+ever written; the rest stays zero-initialized).
9
+
10
+The generic per-group slot kernel cannot express that layout: it maps
11
+``pos -> bt[req][pos // bs] * bs + pos % bs`` (``_compute_slot_mappings_kernel``
12
+in vllm/v1/worker/gpu/block_table.py), so every token at ``pos >= kpool``
13
+reads a zero column and collapses onto physical tail block 0. All concurrent
14
+requests then share one ``kpool``-slot ring and corrupt each other's pool
15
+compression.
16
+
17
+These tests pin that defect's arithmetic, verify the circular replacement
18
+(``compute_kpool_tail_slot_mapping``), and mirror the tail kernels' index
19
+math to show cross-request pollution before the fix and isolation after.
20
+"""
21
+
22
+from types import SimpleNamespace
23
+
24
+import pytest
25
+import torch
26
+
27
+from vllm.v1.attention.backend import CommonAttentionMetadata
28
+from vllm.v1.attention.backends.mla.indexer import (
29
+ KpoolTailBackend,
30
+ KpoolTailMetadataBuilder,
31
+ compute_kpool_tail_slot_mapping,
32
+)
33
+from vllm.v1.kv_cache_interface import KpoolTailSpec, compute_layout_strides
34
+from vllm.v1.kv_cache_layout import KVCacheLayout
35
+
36
+KPOOL = 4
37
+
38
+
39
+def test_tail_backend_layout_matches_kernel_pointer_arithmetic():
40
+ (layout,) = KpoolTailBackend.supported_kv_cache_layouts()
41
+ spec = KpoolTailSpec(
42
+ block_size=KPOOL,
43
+ num_kv_heads=2,
44
+ head_size=128,
45
+ head_size_v=0,
46
+ dtype=torch.bfloat16,
47
+ sliding_window=KPOOL,
48
+ )
49
+ strides = compute_layout_strides(spec, num_blocks=8, num_layers=3, layout=layout)
50
+ _, _, head_stride, state_stride, content_stride = strides
51
+
52
+ assert layout is KVCacheLayout.LBHNC
53
+ assert head_stride == KPOOL * 128 * torch.bfloat16.itemsize
54
+ assert state_stride == 128 * torch.bfloat16.itemsize
55
+ assert content_stride == 1
56
+
57
+
58
+def make_tail_block_table(own_blocks, width=64):
59
+ """Tail-group block table as BlockTables produces it: column 0 holds the
60
+ request's single KpoolTailManager block, the remaining columns are never
61
+ written and stay zero."""
62
+ bt = torch.zeros(len(own_blocks), width, dtype=torch.int32)
63
+ bt[:, 0] = torch.tensor(own_blocks, dtype=torch.int32)
64
+ return bt
65
+
66
+
67
+def legacy_generic_tail_slots(block_table, query_start_loc, positions):
68
+ """Reference of the generic ``_compute_slot_mappings_kernel`` arithmetic
69
+ (block_table.py:305-313) applied to the tail group's table."""
70
+ slots = []
71
+ for req in range(block_table.shape[0]):
72
+ for i in range(query_start_loc[req], query_start_loc[req + 1]):
73
+ pos = int(positions[i])
74
+ block_number = int(block_table[req, pos // KPOOL])
75
+ slots.append(block_number * KPOOL + pos % KPOOL)
76
+ return torch.tensor(slots, dtype=torch.int64)
77
+
78
+
79
+def circular_tail_slots(
80
+ slot_mapping, block_table, query_start_loc, positions, num_actual, num_reqs
81
+):
82
+ return compute_kpool_tail_slot_mapping(
83
+ slot_mapping,
84
+ block_table,
85
+ query_start_loc,
86
+ positions,
87
+ num_actual,
88
+ num_reqs,
89
+ KPOOL,
90
+ )
91
+
92
+
93
+def make_batch(per_req_positions, padded_len=None):
94
+ positions = torch.cat(
95
+ [torch.tensor(p, dtype=torch.int64) for p in per_req_positions]
96
+ )
97
+ num_actual = positions.numel()
98
+ num_reqs = len(per_req_positions)
99
+ lens = [len(p) for p in per_req_positions]
100
+ qsl = torch.zeros(num_reqs + 1, dtype=torch.int64)
101
+ torch.cumsum(torch.tensor(lens, dtype=torch.int64), 0, out=qsl[1:])
102
+ if padded_len is None:
103
+ padded_len = num_actual
104
+ slot_mapping = torch.full((padded_len,), -1, dtype=torch.int64)
105
+ return positions, qsl, slot_mapping, num_actual, num_reqs
106
+
107
+
108
+def test_legacy_generic_mapping_collapses_onto_block_zero():
109
+ """The bug: with the manager's 1-column block table, the generic kernel
110
+ maps every pos >= kpool onto tail block 0, and distinct requests collide."""
111
+ own_blocks = [5, 9]
112
+ per_req = [list(range(10)), list(range(12))] # prompts of len 10 and 12
113
+ positions, qsl, _, num_actual, num_reqs = make_batch(per_req)
114
+ bt = make_tail_block_table(own_blocks)
115
+
116
+ legacy = legacy_generic_tail_slots(bt, qsl, positions)
117
+
118
+ # Every token at pos >= kpool resolves to block 0, not the request's own.
119
+ off = 0
120
+ for req, prompt in enumerate(per_req):
121
+ req_slots = legacy[off : off + len(prompt)]
122
+ for pos in range(len(prompt)):
123
+ slot = int(req_slots[pos])
124
+ if pos >= KPOOL:
125
+ assert slot // KPOOL == 0, (
126
+ f"expected collapse onto block 0 at pos {pos}"
127
+ )
128
+ assert slot // KPOOL != own_blocks[req]
129
+ off += len(prompt)
130
+
131
+ # The two requests share ring slots -> cross-request pollution.
132
+ a_slots = set(legacy[: len(per_req[0])].tolist())
133
+ b_slots = set(legacy[len(per_req[0]) :].tolist())
134
+ assert a_slots & b_slots, "legacy mapping must collide across requests"
135
+
136
+
137
+def test_circular_mapping_isolates_requests():
138
+ """The fix: every token lands in its own request's block at pos % kpool,
139
+ and no slot is ever shared by two different requests (slots do recur
140
+ within a request every kpool positions -- that is the circular design)."""
141
+ own_blocks = [5, 9]
142
+ per_req = [list(range(10)), list(range(12))]
143
+ positions, qsl, slot_mapping, num_actual, num_reqs = make_batch(per_req)
144
+ bt = make_tail_block_table(own_blocks)
145
+
146
+ out = circular_tail_slots(slot_mapping, bt, qsl, positions, num_actual, num_reqs)
147
+
148
+ off = 0
149
+ per_req_slots = []
150
+ for req, prompt in enumerate(per_req):
151
+ req_slots = set()
152
+ for pos in range(len(prompt)):
153
+ slot = int(out[off + pos])
154
+ assert slot // KPOOL == own_blocks[req], (
155
+ f"req {req} pos {pos} left its tail block"
156
+ )
157
+ assert slot % KPOOL == pos % KPOOL
158
+ req_slots.add(slot)
159
+ per_req_slots.append(req_slots)
160
+ off += len(prompt)
161
+ assert not per_req_slots[0] & per_req_slots[1]
162
+
163
+
164
+@pytest.mark.parametrize("prompt_len", [1, 2, 3, 4])
165
+def test_circular_mapping_matches_generic_for_short_requests(prompt_len):
166
+ """For pos < kpool the generic kernel already picks the own block, so the
167
+ two mappings agree while every position fits the request's first block
168
+ (single-request behavior is unchanged)."""
169
+ own_blocks = [7]
170
+ per_req = [list(range(prompt_len))]
171
+ positions, qsl, slot_mapping, num_actual, num_reqs = make_batch(per_req)
172
+ bt = make_tail_block_table(own_blocks)
173
+
174
+ legacy = legacy_generic_tail_slots(bt, qsl, positions)
175
+ out = circular_tail_slots(slot_mapping, bt, qsl, positions, num_actual, num_reqs)
176
+ assert torch.equal(out, legacy)
177
+
178
+
179
+def test_circular_mapping_preserves_padding_and_empty_batch():
180
+ own_blocks = [5, 9]
181
+ per_req = [list(range(10)), list(range(12))]
182
+ padded_len = sum(len(p) for p in per_req) + 8
183
+ positions, qsl, slot_mapping, num_actual, num_reqs = make_batch(
184
+ per_req, padded_len=padded_len
185
+ )
186
+ bt = make_tail_block_table(own_blocks)
187
+
188
+ out = circular_tail_slots(slot_mapping, bt, qsl, positions, num_actual, num_reqs)
189
+ assert out.shape == slot_mapping.shape
190
+ assert torch.equal(out[num_actual:], torch.full_like(out[num_actual:], -1))
191
+
192
+ empty = circular_tail_slots(slot_mapping, bt, qsl, positions[:0], 0, num_reqs)
193
+ assert torch.equal(empty, slot_mapping)
194
+
195
+
196
+def make_common_metadata(per_req_positions, own_blocks, with_positions=True):
197
+ positions, qsl, slot_mapping, num_actual, num_reqs = make_batch(
198
+ per_req_positions, padded_len=sum(len(p) for p in per_req_positions) + 4
199
+ )
200
+ bt = make_tail_block_table(own_blocks)
201
+ seq_lens = torch.tensor(
202
+ [max(p) + 1 if p else 1 for p in per_req_positions], dtype=torch.int64
203
+ )
204
+ return CommonAttentionMetadata(
205
+ query_start_loc=qsl,
206
+ query_start_loc_cpu=qsl.clone(),
207
+ seq_lens=seq_lens,
208
+ num_reqs=num_reqs,
209
+ num_actual_tokens=num_actual,
210
+ max_query_len=max((len(p) for p in per_req_positions), default=1),
211
+ max_seq_len=int(seq_lens.max()) if num_reqs else 1,
212
+ block_table_tensor=bt,
213
+ slot_mapping=slot_mapping,
214
+ positions=positions if with_positions else None,
215
+ )
216
+
217
+
218
+def make_tail_builder(block_size=KPOOL, max_num_batched_tokens=128):
219
+ builder = object.__new__(KpoolTailMetadataBuilder)
220
+ builder.kv_cache_spec = SimpleNamespace(block_size=block_size)
221
+ builder.slot_mapping_buffer = torch.empty(max_num_batched_tokens, dtype=torch.int64)
222
+ return builder
223
+
224
+
225
+def test_builder_build_uses_circular_mapping():
226
+ per_req = [list(range(10)), list(range(12))]
227
+ own_blocks = [5, 9]
228
+ cam = make_common_metadata(per_req, own_blocks)
229
+ meta = KpoolTailMetadataBuilder.build(make_tail_builder(), 0, cam)
230
+
231
+ out = meta.slot_mapping
232
+ off = 0
233
+ for req, prompt in enumerate(per_req):
234
+ for pos in range(len(prompt)):
235
+ slot = int(out[off + pos])
236
+ assert slot // KPOOL == own_blocks[req]
237
+ assert slot % KPOOL == pos % KPOOL
238
+ off += len(prompt)
239
+ # Padding tail of the buffer keeps the -1 sentinel.
240
+ assert torch.equal(
241
+ out[cam.num_actual_tokens :], torch.full_like(out[cam.num_actual_tokens :], -1)
242
+ )
243
+
244
+
245
+def test_builder_build_falls_back_without_positions():
246
+ """Capture / dummy builds without positions keep the generic mapping."""
247
+ per_req = [list(range(10))]
248
+ cam = make_common_metadata(per_req, [5], with_positions=False)
249
+ meta = KpoolTailMetadataBuilder.build(make_tail_builder(), 0, cam)
250
+ assert meta.slot_mapping is cam.slot_mapping
251
+
252
+
253
+def test_builder_reuses_slot_mapping_storage():
254
+ builder = make_tail_builder()
255
+ first = make_common_metadata([list(range(10))], [5])
256
+ first_meta = KpoolTailMetadataBuilder.build(builder, 0, first)
257
+ data_ptr = first_meta.slot_mapping.data_ptr()
258
+
259
+ second = make_common_metadata([list(range(12))], [9])
260
+ second_meta = KpoolTailMetadataBuilder.build(builder, 0, second)
261
+
262
+ assert second_meta.slot_mapping.data_ptr() == data_ptr
263
+ assert second_meta.slot_mapping[:12].tolist() == [
264
+ 9 * KPOOL + pos % KPOOL for pos in range(12)
265
+ ]
266
+
267
+
268
+# ---------------------------------------------------------------------------
269
+# Index-level mirror of the tail kernels: seed / stash / pool completion
270
+# (addressing replicated from kpool_compress.py's Triton kernels).
271
+# ---------------------------------------------------------------------------
272
+
273
+
274
+class TailRingMirror:
275
+ """Mirror of _kpool_tail_seed_kernel / _kpool_decode_update_batched_kernel
276
+ addressing: block = tail_slot // kpool, ring offset = pos % kpool; a pool
277
+ completing at pos reads ring slots (pool_start + s) % kpool and uses the
278
+ current token's own K/score for the last member."""
279
+
280
+ def __init__(self, num_blocks, kpool=KPOOL):
281
+ self.kpool = kpool
282
+ self.k = torch.full((num_blocks, kpool, 3), float("nan"))
283
+ self.s = torch.full((num_blocks, kpool, 3), float("nan"))
284
+
285
+ def stash(self, tail_slot, pos, k, s):
286
+ blk, off = tail_slot // self.kpool, pos % self.kpool
287
+ self.k[blk, off] = k
288
+ self.s[blk, off] = s
289
+
290
+ seed = stash # the seed kernel writes with the same addressing
291
+
292
+ def complete(self, tail_slot, pos, k, s):
293
+ blk = tail_slot // self.kpool
294
+ start = pos - (self.kpool - 1)
295
+ kk = torch.stack(
296
+ [self.k[blk, (start + i) % self.kpool] for i in range(self.kpool)]
297
+ )
298
+ ss = torch.stack(
299
+ [self.s[blk, (start + i) % self.kpool] for i in range(self.kpool)]
300
+ )
301
+ kk[-1], ss[-1] = k, s # is_current for the completing token
302
+ w = torch.softmax(ss, dim=0)
303
+ return (kk * w).sum(0)
304
+
305
+
306
+def token_kv(req, pos):
307
+ k = torch.tensor([pos + 100.0 * req, pos + 0.5, 2.0 * pos + 0.25])
308
+ s = torch.tensor([0.1 * (pos + 1) + req, 0.2 * pos, 0.05 * pos])
309
+ return k, s
310
+
311
+
312
+def tail_slot_for(mapping, req, pos, own_block):
313
+ if mapping == "legacy":
314
+ bt_val = own_block if pos < KPOOL else 0
315
+ return bt_val * KPOOL + pos % KPOOL
316
+ return own_block * KPOOL + pos % KPOOL
317
+
318
+
319
+def run_scenario(mapping, interleave):
320
+ """Requests A (block 5, prompt len 9) and B (block 9, prompt len 11)
321
+ decode concurrently; returns A's boundary pool [8, 9, 10, 11]."""
322
+ ring = TailRingMirror(num_blocks=16)
323
+ blocks = {"A": 5, "B": 9}
324
+ prompts = {"A": 9, "B": 11}
325
+
326
+ def slot(req, pos):
327
+ return tail_slot_for(mapping, 0 if req == "A" else 1, pos, blocks[req])
328
+
329
+ # Prefill: seed each request's trailing incomplete pool.
330
+ for req, L in prompts.items():
331
+ for pos in range(L - (L % KPOOL or KPOOL), L):
332
+ if pos < 0:
333
+ continue
334
+ ring.stash(slot(req, pos), pos, *token_kv(0 if req == "A" else 1, pos))
335
+
336
+ # Decode: A emits pos 9, 10, 11; B emits 11, 12, 13. `interleave`
337
+ # processes B before A within a step, which is exactly what concurrent
338
+ # Triton programs do when they share tail block 0.
339
+ order = ["B", "A"] if interleave else ["A", "B"]
340
+ decode = {"A": [9, 10, 11], "B": [11, 12, 13]}
341
+ result = None
342
+ for step in range(3):
343
+ for req in order:
344
+ pos = decode[req][step]
345
+ k, s = token_kv(0 if req == "A" else 1, pos)
346
+ if pos % KPOOL == KPOOL - 1:
347
+ pool = ring.complete(slot(req, pos), pos, k, s)
348
+ if req == "A":
349
+ result = pool
350
+ ring.stash(slot(req, pos), pos, k, s)
351
+ return result
352
+
353
+
354
+def test_interleaved_decode_pollution_legacy_vs_circular():
355
+ """Ground truth: request A decoding alone (its ring used exclusively)."""
356
+ torch.manual_seed(0)
357
+ ground_truth = run_scenario("circular", interleave=False)
358
+
359
+ legacy = run_scenario("legacy", interleave=True)
360
+ circular = run_scenario("circular", interleave=True)
361
+
362
+ # The old mapping lets request B's tokens into A's ring: A's boundary
363
+ # pool is compressed from 2 of B's tokens -> wrong.
364
+ assert not torch.allclose(legacy, ground_truth), (
365
+ f"legacy mapping unexpectedly clean: {legacy} vs {ground_truth}"
366
+ )
367
+
368
+ # The circular mapping keeps the rings isolated under interleaving.
369
+ torch.testing.assert_close(circular, ground_truth)
tests/v1/attention/test_rocm_glm5next_sparse.py
new
+127
@@ -0,0 +1,127 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+
7
+from vllm.platforms import current_platform
8
+from vllm.triton_utils import tl, triton
9
+from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import (
10
+ _use_rocm_sparse_triton,
11
+ fit_kpool_indices_to_aiter,
12
+)
13
+from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
14
+ _sparse_kv_row_offset,
15
+ _validate_dsv4_sparse_dims,
16
+ _validate_sparse_dims,
17
+)
18
+
19
+
20
+@triton.jit
21
+def _store_sparse_kv_row_offset_kernel(slot_ptr, output_ptr, stride: tl.constexpr):
22
+ slot = tl.load(slot_ptr)
23
+ tl.store(output_ptr, _sparse_kv_row_offset(slot, stride))
24
+
25
+
26
+def test_fit_kpool_indices_preserves_tail_and_best_history():
27
+ token_indices = torch.tensor(
28
+ [
29
+ [10, 9, 8, 7, 6, 5, 100, 101],
30
+ [10, 9, 8, -1, -1, -1, 100, -1],
31
+ [-1, -1, -1, -1, -1, -1, -1, -1],
32
+ ],
33
+ dtype=torch.int32,
34
+ )
35
+
36
+ fitted = fit_kpool_indices_to_aiter(token_indices, topk_tokens=6)
37
+
38
+ assert fitted.tolist() == [
39
+ [10, 9, 8, 7, 100, 101],
40
+ [10, 9, 8, 100, -1, -1],
41
+ [-1, -1, -1, -1, -1, -1],
42
+ ]
43
+
44
+
45
+def test_fit_kpool_indices_exact_width_is_noop():
46
+ token_indices = torch.tensor([[3, 2, 1, -1]], dtype=torch.int32)
47
+
48
+ fitted = fit_kpool_indices_to_aiter(token_indices, topk_tokens=4)
49
+
50
+ assert fitted.data_ptr() == token_indices.data_ptr()
51
+
52
+
53
+def test_fit_kpool_indices_rejects_narrow_input():
54
+ with pytest.raises(ValueError, match="at least topk_tokens"):
55
+ fit_kpool_indices_to_aiter(
56
+ torch.zeros((1, 3), dtype=torch.int32), topk_tokens=4
57
+ )
58
+
59
+
60
+@pytest.mark.parametrize(
61
+ (
62
+ "kv_cache_dtype",
63
+ "head_size",
64
+ "num_prefills",
65
+ "num_decodes",
66
+ "num_decode_tokens",
67
+ "max_query_len",
68
+ "expected",
69
+ ),
70
+ [
71
+ ("auto", 512, 1, 0, 0, 32, True),
72
+ ("auto", 512, 1, 2, 2, 32, True),
73
+ ("auto", 512, 0, 2, 2, 1, True),
74
+ ("fp8", 512, 1, 0, 0, 32, False),
75
+ ("auto", 576, 1, 0, 0, 32, False),
76
+ ("auto", 512, 0, 2, 4, 2, False),
77
+ ],
78
+)
79
+def test_rocm_sparse_triton_route(
80
+ kv_cache_dtype,
81
+ head_size,
82
+ num_prefills,
83
+ num_decodes,
84
+ num_decode_tokens,
85
+ max_query_len,
86
+ expected,
87
+):
88
+ assert (
89
+ _use_rocm_sparse_triton(
90
+ kv_cache_dtype=kv_cache_dtype,
91
+ head_size=head_size,
92
+ kv_lora_rank=512,
93
+ num_prefills=num_prefills,
94
+ num_decodes=num_decodes,
95
+ num_decode_tokens=num_decode_tokens,
96
+ max_query_len=max_query_len,
97
+ )
98
+ is expected
99
+ )
100
+
101
+
102
+def test_rocm_sparse_attention_accepts_glm_nope_dimensions():
103
+ _validate_sparse_dims(512, 512, 0, "test")
104
+
105
+
106
+def test_rocm_sparse_attention_rejects_inconsistent_dimensions():
107
+ with pytest.raises(AssertionError, match="expected head_dim"):
108
+ _validate_sparse_dims(511, 512, 0, "test")
109
+
110
+
111
+def test_dsv4_sparse_attention_keeps_layout_constraint():
112
+ _validate_dsv4_sparse_dims(512, 448, 64, "test")
113
+ with pytest.raises(AssertionError, match="expects 448 NoPE dims"):
114
+ _validate_dsv4_sparse_dims(512, 512, 0, "test")
115
+
116
+
117
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm required")
118
+def test_sparse_prefill_kv_row_offset_does_not_overflow_int32():
119
+ # GLM's 640-token pages cross the signed-int32 address boundary at block
120
+ # 6554 for a 512-element KV row. The production kernel must promote the
121
+ # slot before multiplying by the row stride.
122
+ slot = torch.tensor([6554 * 640], dtype=torch.int32, device="cuda")
123
+ output = torch.empty(1, dtype=torch.int64, device="cuda")
124
+
125
+ _store_sparse_kv_row_offset_kernel[(1,)](slot, output, stride=512)
126
+
127
+ assert output.item() == 6554 * 640 * 512
tests/v1/attention/test_sparse_indexer_decode_seq_lens.py
new
+184
@@ -0,0 +1,184 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""CPU tests for the decode top-k row->token mapping (no GPU required).
4
+
5
+On non-uniform decode batches (``requires_padding`` -- mixed plain-decode and
6
+spec-verify requests, or variable MTP verify lens; taken on Hopper where the
7
+varlen/flatten logits path is unavailable), the pool-topk rows follow the
8
+PADDED ``[batch_size, next_n]`` grid: row ``(b, t)`` is flat decode token
9
+``offset_b + t``. The former inline ``dec_seq = positions[:n] + 1``
10
+(``n = batch_size * next_n``) indexes the flat per-token layout with padded
11
+coordinates, so rows after the first non-uniform request read another
12
+request's positions and rows past the decode region read prefill tokens --
13
+``expand_pools_and_append_tail`` then anchors the tail at a foreign length.
14
+
15
+These tests pin the defect with the exact production arithmetic and verify
16
+the layout-aware replacement (``_decode_topk_seq_lens``), including the tail
17
+expansion consequences via the pure-torch expand/append pair that the fused
18
+kernel is documented to replicate.
19
+"""
20
+
21
+import torch
22
+
23
+# Bootstrap the glm5next package before entering the indexer module: its
24
+# kpool_compress import runs glm5next/__init__, which pulls model ->
25
+# attention -> back into sparse_attn_indexer_kpool (attention.py imports the
26
+# class at module scope). Production always enters via attention.py first.
27
+# isort: off
28
+import vllm.models.glm5next # noqa: F401
29
+from vllm.models.glm5next.nvidia.ops.kpool_compress import ( # noqa: E402
30
+ append_tail_to_topk,
31
+ expand_pools_to_tokens,
32
+)
33
+# isort: on
34
+
35
+import vllm.model_executor.layers.sparse_attn_indexer_kpool as indexer_mod
36
+from vllm.model_executor.layers.sparse_attn_indexer_kpool import (
37
+ _decode_topk_seq_lens,
38
+ _fill_short_decode_causal_indices,
39
+)
40
+from vllm.platforms import current_platform
41
+
42
+KPOOL = 4
43
+TOPK_TOKENS = 16
44
+SELECT_K = TOPK_TOKENS // KPOOL
45
+
46
+
47
+def test_kpool_ops_dispatch_matches_platform():
48
+ expected_backend = ".amd." if current_platform.is_rocm() else ".nvidia."
49
+ assert expected_backend in indexer_mod.kpool_ops.__name__
50
+
51
+
52
+def test_short_decode_fills_exact_causal_rows():
53
+ topk = torch.full((3, 8), 99, dtype=torch.int32)
54
+ positions = torch.tensor([0, 3, 7], dtype=torch.int64)
55
+
56
+ assert _fill_short_decode_causal_indices(topk, positions, 3, 8, 8)
57
+ assert topk.tolist() == [
58
+ [0, -1, -1, -1, -1, -1, -1, -1],
59
+ [0, 1, 2, 3, -1, -1, -1, -1],
60
+ [0, 1, 2, 3, 4, 5, 6, 7],
61
+ ]
62
+
63
+
64
+def test_short_decode_leaves_buffer_unchanged_for_sparse_context():
65
+ topk = torch.full((2, 8), 99, dtype=torch.int32)
66
+ before = topk.clone()
67
+
68
+ assert not _fill_short_decode_causal_indices(topk, torch.tensor([7, 8]), 2, 9, 8)
69
+ assert torch.equal(topk, before)
70
+
71
+
72
+def make_non_uniform_batch():
73
+ """3 requests: plain decode (1 token), MTP verify (4), adaptive verify (3).
74
+
75
+ Flat decode positions (production layout: decode tokens first, then any
76
+ prefill tokens of the same batch):
77
+ req0: [30] (context len 30)
78
+ req1: [100..103] (verify at context len 100)
79
+ req2: [7, 8, 9] (verify at context len 7)
80
+ Followed by 4 prefill tokens at positions 555..558 so the flat tensor is
81
+ at least ``n = 3 * 4`` long, as it is in a real mixed batch.
82
+ """
83
+ decode_lens = torch.tensor([1, 4, 3], dtype=torch.int64)
84
+ per_req_positions = [[30], [100, 101, 102, 103], [7, 8, 9]]
85
+ flat_decode = [p for req in per_req_positions for p in req]
86
+ positions = torch.tensor(flat_decode + [555, 556, 557, 558], dtype=torch.int64)
87
+ return decode_lens, per_req_positions, positions
88
+
89
+
90
+def expected_row_seq_lens(per_req_positions, batch_size, next_n):
91
+ """Ground truth: row (b, t) -> pos + 1 for real rows, 0 for pad rows."""
92
+ out = torch.zeros(batch_size * next_n, dtype=torch.int32)
93
+ for b, req in enumerate(per_req_positions):
94
+ for t, pos in enumerate(req):
95
+ out[b * next_n + t] = pos + 1
96
+ return out
97
+
98
+
99
+def test_legacy_flat_layout_misaligns_and_bleeds():
100
+ """The bug: ``positions[:n] + 1`` with padded-row coordinates reads other
101
+ requests' positions and prefill positions."""
102
+ decode_lens, per_req, positions = make_non_uniform_batch()
103
+ batch_size = decode_lens.shape[0]
104
+ next_n = int(decode_lens.max())
105
+ n = batch_size * next_n
106
+ assert n == 12
107
+
108
+ legacy = positions[:n].to(torch.int32) + 1
109
+ expected = expected_row_seq_lens(per_req, batch_size, next_n)
110
+
111
+ # Row (1, 0): req1's first verify token (pos 100, seq 101) reads flat
112
+ # index 4 -> req1's LAST verify token (pos 103). Its true tail token is
113
+ # dropped and the tail anchors 3 tokens late.
114
+ assert int(legacy[4]) == 104 and int(expected[4]) == 101
115
+
116
+ # Rows (2, *): req2 starts at flat offset 5, but padded coordinates point
117
+ # at flat indices 8..11 -- the batch's PREFILL positions.
118
+ for t in range(3):
119
+ assert int(legacy[2 * next_n + t]) >= 556, (
120
+ "legacy row (2, t) should read prefill positions"
121
+ )
122
+ assert not torch.equal(legacy, expected)
123
+
124
+
125
+def test_helper_uniform_layout_matches_flat_slice():
126
+ """Uniform batches keep the flat shortcut (zero behavior/perf change)."""
127
+ per_req = [[200, 201, 202, 203], [50, 51, 52, 53]]
128
+ positions = torch.tensor([p for r in per_req for p in r], dtype=torch.int64)
129
+ decode_lens = torch.tensor([4, 4], dtype=torch.int64)
130
+ out = _decode_topk_seq_lens(positions, decode_lens, 8, 2, 4, requires_padding=False)
131
+ assert torch.equal(out, positions[:8].to(torch.int32) + 1)
132
+
133
+
134
+def test_helper_padded_layout_per_row():
135
+ """Non-uniform batches map every padded row to its own token's position;
136
+ pad rows get 0 (empty tail)."""
137
+ decode_lens, per_req, positions = make_non_uniform_batch()
138
+ out = _decode_topk_seq_lens(positions, decode_lens, 8, 3, 4, requires_padding=True)
139
+ expected = expected_row_seq_lens(per_req, 3, 4)
140
+ assert torch.equal(out, expected)
141
+ # Pad rows (0, 1..3) and (2, 3) collapse to 0 -> no tail appended.
142
+ assert out[1] == 0 and out[2] == 0 and out[3] == 0 and out[11] == 0
143
+
144
+
145
+def expand_tail_region(dec_seq):
146
+ """Run the production pure-torch expand + append pair (the fused kernel
147
+ replicates it exactly on the identity path) and return the tail columns
148
+ [TOPK_TOKENS, TOPK_TOKENS + KPOOL - 1)."""
149
+ rows = dec_seq.shape[0]
150
+ pool_ids = torch.arange(SELECT_K, dtype=torch.int64).expand(rows, SELECT_K)
151
+ valid = torch.ones_like(pool_ids, dtype=torch.bool)
152
+ expanded = expand_pools_to_tokens(pool_ids, valid, TOPK_TOKENS, KPOOL)
153
+ seq_lens = dec_seq.to(torch.int32)
154
+ pool_lens = (seq_lens // KPOOL).to(torch.int32)
155
+ out = append_tail_to_topk(expanded, seq_lens, pool_lens, KPOOL)
156
+ return out[:, TOPK_TOKENS:]
157
+
158
+
159
+def test_tail_expansion_legacy_vs_fixed():
160
+ """End-to-end tail consequence: the fixed mapping appends exactly the
161
+ request's trailing incomplete pool; the legacy flat slice drops real tail
162
+ tokens and emits indices far past the request's own sequence."""
163
+ decode_lens, per_req, positions = make_non_uniform_batch()
164
+ n = 12
165
+ fixed = _decode_topk_seq_lens(
166
+ positions, decode_lens, 8, 3, 4, requires_padding=True
167
+ )
168
+ legacy = positions[:n].to(torch.int32) + 1
169
+
170
+ fixed_tail = expand_tail_region(fixed)
171
+ legacy_tail = expand_tail_region(legacy)
172
+
173
+ # Fixed: row (1, 0) (seq 101) keeps its single tail token 100; row (2, 2)
174
+ # (seq 10) keeps its full trailing pool [8, 9].
175
+ assert fixed_tail[4, 0].item() == 100
176
+ assert fixed_tail[10, :2].tolist() == [8, 9]
177
+
178
+ # Legacy: row (1, 0) loses its tail token (anchored at 104, count 0)...
179
+ assert legacy_tail[4, 0].item() == -1
180
+ # ...and row (2, 2) reads prefill position 557 -> tail indices 556/557,
181
+ # way past req2's 10-token sequence -> out-of-bounds block-table reads.
182
+ assert legacy_tail[10, :2].tolist() == [556, 557]
183
+
184
+ assert not torch.equal(legacy_tail, fixed_tail)
tests/v1/attention/test_sparse_mla_backends.py
+61
-5
@@ -40,11 +40,13 @@ if not current_platform.is_cuda():
40
allow_module_level=True,
41
)
42
43
+import vllm.v1.attention.backends.mla.flashinfer_mla_sparse as flashinfer_sparse_mod
44
from vllm.model_executor.layers.attention.mla_attention import (
45
_canonicalize_sparse_mla_kv_cache_dtype,
46
)
47
from vllm.utils.math_utils import cdiv
48
from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
49
+ FlashInferMLASparseImpl,
50
FlashInferMLASparseTRTLLMBackend,
51
)
52
from vllm.v1.attention.backends.mla.flashmla_sparse import (
@@ -82,6 +84,60 @@ SPARSE_BACKEND_BATCH_SPECS["large_q_pure_prefill"] = BatchSpec(
84
DEVICE_TYPE = current_platform.device_type
85
86
87
+def test_nope_flashinfer_sparse_mla_uses_model_scale(monkeypatch):
88
+ """Weight absorption must not change the model's attention temperature."""
89
+ model_scale = 256**-0.5
90
+ kv_lora_rank = 512
91
+ topk = torch.zeros((1, 1), dtype=torch.int32)
92
+ metadata = SimpleNamespace(
93
+ req_id_per_token=torch.zeros(1, dtype=torch.int32),
94
+ block_table=torch.zeros((1, 1), dtype=torch.int32),
95
+ block_size=1,
96
+ )
97
+ recorded_scale = None
98
+
99
+ impl = object.__new__(FlashInferMLASparseImpl)
100
+ impl.scale = model_scale
101
+ impl.qk_nope_head_dim = 256
102
+ impl.kv_lora_rank = kv_lora_rank
103
+ impl.qk_rope_head_dim = 0
104
+ impl.kv_cache_dtype = "auto"
105
+ impl.topk_indices_buffer = topk
106
+ impl.dcp_world_size = 1
107
+ impl._workspace_buffer = torch.empty(1)
108
+ impl.bmm1_scale = None
109
+ impl.bmm2_scale = None
110
+ impl.is_nope_mla = True
111
+ impl.need_to_return_lse_for_decode = False
112
+ monkeypatch.setattr(
113
+ flashinfer_sparse_mod,
114
+ "triton_convert_req_index_to_global_index",
115
+ lambda *args, **kwargs: (topk, torch.ones(1, dtype=torch.int32)),
116
+ )
117
+
118
+ import flashinfer.decode
119
+
120
+ def fake_flashinfer(**kwargs):
121
+ nonlocal recorded_scale
122
+ recorded_scale = kwargs["bmm1_scale"]
123
+ return torch.zeros((1, 1, 1, kv_lora_rank))
124
+
125
+ monkeypatch.setattr(
126
+ flashinfer.decode,
127
+ "trtllm_batch_decode_with_kv_cache_mla",
128
+ fake_flashinfer,
129
+ )
130
+ impl.forward_mqa(
131
+ torch.zeros(1, 1, kv_lora_rank),
132
+ torch.zeros(1, 1, kv_lora_rank),
133
+ metadata,
134
+ SimpleNamespace(),
135
+ )
136
+
137
+ assert recorded_scale == model_scale
138
+ assert recorded_scale != kv_lora_rank**-0.5
139
+
140
+
141
def _float_to_e8m0_truncate(f: float) -> float:
142
"""Simulate SM100's float -> e8m0 -> bf16 scale conversion.
143
e8m0 format only stores the exponent (power of 2).
@@ -1451,14 +1507,14 @@ def test_triton_convert_returns_valid_counts(num_topk_tokens: int):
1507
return_valid_counts=False,
1508
)
1509
assert isinstance(result_only, torch.Tensor)
1454
- torch.testing.assert_close(result_only, result, rtol=0, atol=0)
1510
+ for row, num_valid in enumerate(expected_valid):
1511
+ compact_valid = result[row, :num_valid].sort().values
1512
+ original_valid = result_only[row][result_only[row] >= 0].sort().values
1513
+ torch.testing.assert_close(compact_valid, original_valid, rtol=0, atol=0)
1514
+ assert torch.all(result[row, num_valid:] == -1)
1515
1516
1517
def test_flashmla_cache_dtype_aliases_use_ds_layout():
1458
- from vllm.model_executor.layers.attention.mla_attention import (
1459
- _canonicalize_sparse_mla_kv_cache_dtype,
1460
- )
1461
-
1518
# kv-cache dtype aliases are canonicalized to fp8_ds_mla before the layer
1519
# stores kv_cache_dtype, so they cannot bypass the gate.
1520
for alias in ("fp8", "fp8_e4m3"):
tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py
+122
@@ -23,6 +23,7 @@ from vllm.v1.core.kv_cache_utils import (
23
from vllm.v1.core.sched.scheduler import Scheduler
24
from vllm.v1.kv_cache_interface import (
25
FullAttentionSpec,
26
+ KpoolTailSpec,
27
KVCacheConfig,
28
KVCacheGroupSpec,
29
MambaSpec,
@@ -1801,6 +1802,127 @@ def test_hybrid_sliding_window_group_disables_partial_hash_hits():
1802
assert len(computed_blocks.blocks[0]) * hash_block_size == num_computed
1803
1804
1805
+def test_opted_out_scratch_group_keeps_partial_hash_hits():
1806
+ hash_block_size = 2
1807
+ mamba_block_size = 2 * hash_block_size
1808
+ kv_cache_config = KVCacheConfig(
1809
+ num_blocks=24,
1810
+ kv_cache_tensors=[],
1811
+ kv_cache_groups=[
1812
+ KVCacheGroupSpec(
1813
+ ["full"],
1814
+ FullAttentionSpec(
1815
+ block_size=hash_block_size,
1816
+ num_kv_heads=1,
1817
+ head_size=1,
1818
+ dtype=torch.float32,
1819
+ ),
1820
+ ),
1821
+ KVCacheGroupSpec(
1822
+ ["mamba"],
1823
+ MambaSpec(
1824
+ block_size=mamba_block_size,
1825
+ shapes=(1, 1),
1826
+ dtypes=(torch.float32,),
1827
+ mamba_cache_mode="align",
1828
+ ),
1829
+ ),
1830
+ KVCacheGroupSpec(
1831
+ ["tail"],
1832
+ KpoolTailSpec(
1833
+ block_size=mamba_block_size,
1834
+ num_kv_heads=1,
1835
+ head_size=1,
1836
+ dtype=torch.float32,
1837
+ sliding_window=mamba_block_size,
1838
+ ),
1839
+ ),
1840
+ ],
1841
+ )
1842
+ manager = make_kv_cache_manager(
1843
+ kv_cache_config=kv_cache_config,
1844
+ max_model_len=8192,
1845
+ enable_caching=True,
1846
+ hash_block_size=hash_block_size,
1847
+ )
1848
+
1849
+ assert manager.coordinator.enable_partial_hash_hits
1850
+
1851
+ req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
1852
+ computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)
1853
+ assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None
1854
+ manager.free(req0)
1855
+ manager.new_step_starts()
1856
+
1857
+ req1 = make_request("1", [0, 0, 1, 1, 2, 2, 3, 3], hash_block_size, sha256)
1858
+ _, num_computed, _ = manager.get_computed_blocks(req1)
1859
+ assert num_computed == 6
1860
+
1861
+
1862
+def test_kpool_tail_supports_128_token_partial_hash_hits():
1863
+ hash_block_size = 128
1864
+ cache_block_size = 9 * hash_block_size
1865
+ kpool_block_size = 64
1866
+ kv_cache_config = KVCacheConfig(
1867
+ num_blocks=24,
1868
+ kv_cache_tensors=[],
1869
+ kv_cache_groups=[
1870
+ KVCacheGroupSpec(
1871
+ ["full"],
1872
+ FullAttentionSpec(
1873
+ block_size=cache_block_size,
1874
+ num_kv_heads=1,
1875
+ head_size=1,
1876
+ dtype=torch.float32,
1877
+ ),
1878
+ ),
1879
+ KVCacheGroupSpec(
1880
+ ["mamba"],
1881
+ MambaSpec(
1882
+ block_size=cache_block_size,
1883
+ shapes=(1, 1),
1884
+ dtypes=(torch.float32,),
1885
+ mamba_cache_mode="align",
1886
+ ),
1887
+ ),
1888
+ KVCacheGroupSpec(
1889
+ ["tail"],
1890
+ KpoolTailSpec(
1891
+ block_size=kpool_block_size,
1892
+ num_kv_heads=1,
1893
+ head_size=1,
1894
+ dtype=torch.float32,
1895
+ sliding_window=kpool_block_size,
1896
+ ),
1897
+ ),
1898
+ ],
1899
+ )
1900
+ manager = make_kv_cache_manager(
1901
+ kv_cache_config=kv_cache_config,
1902
+ max_model_len=8192,
1903
+ enable_caching=True,
1904
+ hash_block_size=hash_block_size,
1905
+ )
1906
+
1907
+ assert manager.coordinator.scheduler_block_size == cache_block_size
1908
+ assert manager.coordinator.hash_block_size == hash_block_size
1909
+ assert manager.coordinator.enable_partial_hash_hits
1910
+
1911
+ shared_prefix = [10] * (12 * hash_block_size)
1912
+ req0 = make_request("0", shared_prefix, hash_block_size, sha256)
1913
+ computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)
1914
+ assert (
1915
+ manager.allocate_slots(req0, len(shared_prefix), num_computed, computed_blocks)
1916
+ is not None
1917
+ )
1918
+ manager.free(req0)
1919
+ manager.new_step_starts()
1920
+
1921
+ req1 = make_request("1", shared_prefix + [12] * 64, hash_block_size, sha256)
1922
+ _, num_computed, _ = manager.get_computed_blocks(req1)
1923
+ assert num_computed == len(shared_prefix)
1924
+
1925
+
1926
@pytest.mark.parametrize("dcp_world_size", [1, 2, 4])
1927
def test_hybrid_partial_hash_hit_uses_cow_under_dcp(dcp_world_size: int):
1928
hash_block_size = 2
tests/v1/core/test_kv_cache_utils.py
+569
-28
@@ -6,13 +6,19 @@ import importlib
6
from collections.abc import Callable
7
from dataclasses import replace
8
from types import SimpleNamespace
9
-from typing import Any
9
+from typing import Any, cast
10
11
import pytest
12
import torch
13
14
import vllm.v1.core.kv_cache_utils as kv_cache_utils
15
-from vllm.config import CacheConfig, ModelConfig, SchedulerConfig, VllmConfig
15
+from vllm.config import (
16
+ CacheConfig,
17
+ KVTransferConfig,
18
+ ModelConfig,
19
+ SchedulerConfig,
20
+ VllmConfig,
21
+)
22
from vllm.config.kv_events import KVEventsConfig
23
from vllm.lora.request import LoRARequest
24
from vllm.multimodal.inputs import (
@@ -47,6 +53,7 @@ from vllm.v1.kv_cache_interface import (
53
ChunkedLocalAttentionSpec,
54
FullAttentionSpec,
55
HiddenStateCacheSpec,
56
+ KpoolTailSpec,
57
KVCacheConfig,
58
KVCacheGroupSpec,
59
KVCacheSpec,
@@ -2191,34 +2198,567 @@ def test_generate_scheduler_kv_cache_config():
2198
)
2199
2200
2194
-def test_mixed_precision_kv_cache_with_uniform_type_specs():
2195
- fp8_spec = new_kv_cache_spec(dtype=torch.float8_e4m3fn)
2196
- bf16_spec = new_kv_cache_spec(dtype=torch.bfloat16)
2197
- worker_config = KVCacheConfig(
2198
- num_blocks=10,
2201
+def _glm5_like_kv_cache_spec(
2202
+ mamba_spec_factory=new_mamba_spec,
2203
+) -> tuple[dict[str, KVCacheSpec], list[str]]:
2204
+ """(mamba, mamba, mamba, MLA + indexer) * 11 plus a trailing mamba."""
2205
+ kv_cache_spec: dict[str, KVCacheSpec] = {}
2206
+ mamba_layers = []
2207
+ for i in range(45):
2208
+ if i % 4 == 3:
2209
+ kv_cache_spec[f"layers.{i}.attn"] = MLAAttentionSpec(
2210
+ block_size=1024,
2211
+ num_kv_heads=1,
2212
+ head_size=576,
2213
+ dtype=torch.bfloat16,
2214
+ )
2215
+ kv_cache_spec[f"layers.{i}.indexer"] = MLAAttentionSpec(
2216
+ block_size=1024,
2217
+ num_kv_heads=1,
2218
+ head_size=132,
2219
+ dtype=torch.uint8,
2220
+ tokens_per_state=16,
2221
+ )
2222
+ else:
2223
+ name = f"layers.{i}.linear_attn"
2224
+ kv_cache_spec[name] = mamba_spec_factory()
2225
+ mamba_layers.append(name)
2226
+ return kv_cache_spec, mamba_layers
2227
+
2228
+
2229
+def _glm5_like_kv_cache_spec_with_tail(
2230
+ mamba_spec_factory=new_mamba_spec,
2231
+) -> dict[str, KVCacheSpec]:
2232
+ """Production kpool=4 proportions: (mamba*3, MLA + indexer + tail) * 11.
2233
+
2234
+ The tail's logical page (2 * kpool * 2*indexer_head_dim * bf16) must fit
2235
+ inside the indexer page it parasitizes (block_size//kpool * 132 B), which
2236
+ is why the fixture drops the tokens_per_state=16 shape used above.
2237
+ """
2238
+ kv_cache_spec, _ = _glm5_like_kv_cache_spec(mamba_spec_factory)
2239
+ for i in range(3, 45, 4):
2240
+ kpool = 4
2241
+ kv_cache_spec[f"layers.{i}.indexer"] = replace(
2242
+ cast(MLAAttentionSpec, kv_cache_spec[f"layers.{i}.indexer"]),
2243
+ tokens_per_state=kpool,
2244
+ )
2245
+ kv_cache_spec[f"layers.{i}.tail"] = KpoolTailSpec(
2246
+ block_size=kpool,
2247
+ num_kv_heads=2,
2248
+ head_size=128,
2249
+ head_size_v=0,
2250
+ dtype=torch.bfloat16,
2251
+ sliding_window=kpool,
2252
+ )
2253
+ return kv_cache_spec
2254
+
2255
+
2256
+def _tensor_by_layer(kv_cache_config: KVCacheConfig) -> dict[str, KVCacheTensor]:
2257
+ return {
2258
+ layer_name: tensor
2259
+ for tensor in kv_cache_config.kv_cache_tensors
2260
+ for layer_name in tensor.layers
2261
+ }
2262
+
2263
+
2264
+def _layer_offset(tensor: KVCacheTensor, layer_name: str) -> int:
2265
+ return tensor.offset + tensor.layers.index(layer_name) * tensor.layer_stride
2266
+
2267
+
2268
+def test_get_kv_cache_config_balanced_mamba_hybrid():
2269
+ """Hybrid slot sharing: mamba layers co-own the MLA slot tensors."""
2270
+ model_config = ModelConfig(max_model_len=8192)
2271
+ vllm_config = VllmConfig(model_config=model_config)
2272
+
2273
+ kv_cache_spec, mamba_layers = _glm5_like_kv_cache_spec()
2274
+ mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes
2275
+ idx_page = kv_cache_spec["layers.3.indexer"].page_size_bytes
2276
+
2277
+ groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec)
2278
+ uniform_groups = [
2279
+ group
2280
+ for group in groups
2281
+ if isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs)
2282
+ ]
2283
+ mamba_groups = [
2284
+ group for group in groups if isinstance(group.kv_cache_spec, MambaSpec)
2285
+ ]
2286
+
2287
+ # Round-robin into G = ceil(34 / 11) = 4 groups so every mamba layer
2288
+ # gets an MLA slot.
2289
+ assert len(groups) == 5
2290
+ assert len(uniform_groups) == 1
2291
+ assert len(uniform_groups[0].layer_names) == 22
2292
+ assert uniform_groups[0].kv_cache_spec.get_max_layers_per_page_size() == 11
2293
+ assert [len(group.layer_names) for group in mamba_groups] == [9, 9, 8, 8]
2294
+ for k, name in enumerate(mamba_layers):
2295
+ assert name in mamba_groups[k % 4].layer_names
2296
+ # Mamba pages are padded up to the MLA page.
2297
+ for group in mamba_groups:
2298
+ assert group.kv_cache_spec.page_size_padded == mla_page
2299
+ assert group.kv_cache_spec.page_size_bytes == mla_page
2300
+
2301
+ bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups)
2302
+ assert bytes_per_block == 11 * mla_page + 11 * idx_page
2303
+
2304
+ # Every block id is charged the full per-block sum.
2305
+ attn_blocks = uniform_groups[0].kv_cache_spec.max_memory_usage_pages(vllm_config)
2306
+ mamba_blocks_per_group = 1 + new_mamba_spec().num_speculative_blocks
2307
+ blocks_per_request = attn_blocks + 4 * mamba_blocks_per_group
2308
+ assert (
2309
+ kv_cache_utils._max_memory_usage_bytes_from_groups(vllm_config, groups)
2310
+ == blocks_per_request * bytes_per_block
2311
+ )
2312
+
2313
+ available_memory = bytes_per_block * 100 + 1
2314
+ kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups(
2315
+ vllm_config, groups, available_memory
2316
+ )
2317
+ assert kv_cache_config.num_blocks == 100
2318
+
2319
+ # Every logical layer has a view into one backing allocation. Mamba views
2320
+ # alias their corresponding MLA slot by using the same byte offset.
2321
+ assert len(kv_cache_config.kv_cache_tensors) == 56
2322
+ assert {t.size for t in kv_cache_config.kv_cache_tensors} == {bytes_per_block * 100}
2323
+ tensors = _tensor_by_layer(kv_cache_config)
2324
+ mla_layer_names = [f"layers.{i}.attn" for i in range(3, 45, 4)]
2325
+ for i, mla_name in enumerate(mla_layer_names):
2326
+ mla_tensor = tensors[mla_name]
2327
+ assert mla_tensor.block_stride == mla_page
2328
+ for group in mamba_groups:
2329
+ if i < len(group.layer_names):
2330
+ assert tensors[group.layer_names[i]].offset == mla_tensor.offset
2331
+ for i in range(11):
2332
+ assert tensors[f"layers.{4 * i + 3}.indexer"].block_stride == idx_page
2333
+
2334
+ total_allocated = next(iter({t.size for t in kv_cache_config.kv_cache_tensors}))
2335
+ assert total_allocated == bytes_per_block * 100
2336
+ assert 0 <= available_memory - total_allocated < bytes_per_block
2337
+
2338
+ assert get_max_concurrency_for_kv_cache_config(
2339
+ vllm_config, kv_cache_config
2340
+ ) == pytest.approx(100 / blocks_per_request)
2341
+
2342
+
2343
+def test_get_kv_cache_config_kpool_tail_coowns_indexer_tensor():
2344
+ """The kpool tail parasitizes the indexer tensors instead of getting its
2345
+ own: sibling idx/tail tensors paired by layer order, zero standalone tail
2346
+ bytes, one shared block per request, and no prefix-caching leakage from
2347
+ the kpool-sized scratch group."""
2348
+ model_config = ModelConfig(max_model_len=8192)
2349
+ vllm_config = VllmConfig(model_config=model_config)
2350
+
2351
+ kv_cache_spec = _glm5_like_kv_cache_spec_with_tail()
2352
+ mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes
2353
+ idx_page = kv_cache_spec["layers.3.indexer"].page_size_bytes
2354
+ tail_logical_page = kv_cache_spec["layers.3.tail"].page_size_bytes
2355
+ assert tail_logical_page < idx_page
2356
+
2357
+ groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec)
2358
+ tail_group = next(
2359
+ group
2360
+ for group in groups
2361
+ if isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs)
2362
+ and all(
2363
+ isinstance(spec, KpoolTailSpec)
2364
+ for spec in group.kv_cache_spec.kv_cache_specs.values()
2365
+ )
2366
+ )
2367
+ # The tail never prefix-caches; its page is padded up to the indexer page
2368
+ # so the runner's strided view rides the indexer storage.
2369
+ assert not tail_group.kv_cache_spec.prefix_cacheable
2370
+ tail_inner = cast(
2371
+ KpoolTailSpec, tail_group.kv_cache_spec.kv_cache_specs["layers.3.tail"]
2372
+ )
2373
+ assert tail_inner.page_size_padded == idx_page
2374
+
2375
+ bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups)
2376
+ assert bytes_per_block == 11 * mla_page + 11 * idx_page
2377
+
2378
+ kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups(
2379
+ vllm_config, groups, bytes_per_block * 100 + 1
2380
+ )
2381
+ assert kv_cache_config.num_blocks == 100
2382
+ # Tail layers get logical views but no additional storage: each view aliases
2383
+ # its sibling indexer at the same offset in the shared allocation.
2384
+ assert len(kv_cache_config.kv_cache_tensors) == 67
2385
+ tensors = _tensor_by_layer(kv_cache_config)
2386
+ for i in range(11):
2387
+ idx_name = f"layers.{4 * i + 3}.indexer"
2388
+ tail_name = f"layers.{4 * i + 3}.tail"
2389
+ assert tensors[idx_name].offset == tensors[tail_name].offset
2390
+ assert {t.size for t in kv_cache_config.kv_cache_tensors} == {bytes_per_block * 100}
2391
+
2392
+ # The layout detector surfaces the sibling names in layer order for the
2393
+ # accounting and connector paths.
2394
+ layout = kv_cache_utils._glm5_next_tensor_layout(kv_cache_config.kv_cache_groups)
2395
+ assert layout is not None
2396
+ assert layout[3] == [f"layers.{4 * i + 3}.indexer" for i in range(11)]
2397
+ assert layout[6] == [f"layers.{4 * i + 3}.tail" for i in range(11)]
2398
+
2399
+ # Accounting charges exactly one extra shared block per request for the
2400
+ # tail, not a per-sequence allocation.
2401
+ attn_group = next(
2402
+ group
2403
+ for group in groups
2404
+ if isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs)
2405
+ and group is not tail_group
2406
+ )
2407
+ attn_blocks = attn_group.kv_cache_spec.max_memory_usage_pages(vllm_config)
2408
+ mamba_blocks_per_group = 1 + new_mamba_spec().num_speculative_blocks
2409
+ blocks_per_request = attn_blocks + 4 * mamba_blocks_per_group + 1
2410
+ assert (
2411
+ kv_cache_utils._max_memory_usage_bytes_from_groups(vllm_config, groups)
2412
+ == blocks_per_request * bytes_per_block
2413
+ )
2414
+
2415
+
2416
+def test_glm5_kpool_tail_does_not_drag_hash_block_size():
2417
+ """The tail's kpool-sized scratch block (4 tokens) must not constrain the
2418
+ prefix-cache hash granularity: participating groups alone decide it."""
2419
+ model_config = ModelConfig(max_model_len=8192)
2420
+ vllm_config = VllmConfig(model_config=model_config)
2421
+
2422
+ def align_mamba():
2423
+ return new_mamba_spec(mamba_cache_mode="align")
2424
+
2425
+ groups = kv_cache_utils.get_kv_cache_groups(
2426
+ vllm_config, _glm5_like_kv_cache_spec_with_tail(align_mamba)
2427
+ )
2428
+ kv_cache_config = KVCacheConfig(
2429
+ num_blocks=1,
2430
kv_cache_tensors=[],
2200
- kv_cache_groups=[
2201
- KVCacheGroupSpec(
2202
- ["fp8_layer"],
2203
- UniformTypeKVCacheSpecs(
2204
- block_size=16, kv_cache_specs={"fp8_layer": fp8_spec}
2205
- ),
2206
- ),
2207
- KVCacheGroupSpec(
2208
- ["bf16_layer"],
2209
- UniformTypeKVCacheSpecs(
2210
- block_size=16, kv_cache_specs={"bf16_layer": bf16_spec}
2211
- ),
2212
- ),
2213
- ],
2431
+ kv_cache_groups=groups,
2432
+ )
2433
+ hash_vllm_config = SimpleNamespace(
2434
+ cache_config=SimpleNamespace(
2435
+ block_size=16,
2436
+ enable_prefix_caching=True,
2437
+ prefix_match_unit=None,
2438
+ ),
2439
+ parallel_config=SimpleNamespace(decode_context_parallel_size=1),
2440
+ kv_transfer_config=object(),
2441
+ )
2442
+ # gcd(attn 1024, mamba 16) with the tail's 4 excluded; scheduler size is
2443
+ # the lcm including the tail (1024 % 4 == 0, so it coincides).
2444
+ assert kv_cache_utils.resolve_kv_cache_block_sizes(
2445
+ kv_cache_config, hash_vllm_config
2446
+ ) == (1024, 16)
2447
+
2448
+
2449
+def test_get_kv_cache_config_mamba_hybrid_sharing_infeasible():
2450
+ """Reject GLM-5.3-Flash layouts whose Mamba page exceeds the MLA page."""
2451
+ model_config = ModelConfig(max_model_len=8192)
2452
+ vllm_config = VllmConfig(model_config=model_config)
2453
+
2454
+ # (512, 1024) fp32 state = 2 MiB per page > 1.125 MiB MLA page.
2455
+ def big_mamba_spec():
2456
+ return new_mamba_spec(shapes=((512, 1024),), dtypes=(torch.float32,))
2457
+
2458
+ kv_cache_spec, _ = _glm5_like_kv_cache_spec(big_mamba_spec)
2459
+ mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes
2460
+ assert big_mamba_spec().page_size_bytes > mla_page
2461
+
2462
+ with pytest.raises(ValueError, match="does not fit the MLA page"):
2463
+ kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec)
2464
+
2465
+
2466
+def test_get_kv_cache_config_mamba_hybrid_sharing_infeasible_no_indexer():
2467
+ """Use the generic layout error when no kpool indexer is present."""
2468
+ model_config = ModelConfig(max_model_len=8192)
2469
+ vllm_config = VllmConfig(model_config=model_config)
2470
+ vllm_config.cache_config.kv_cache_layout = "LBHNC"
2471
+
2472
+ kv_cache_spec: dict[str, KVCacheSpec] = {}
2473
+ for i in range(27):
2474
+ if i % 4 == 3 or i == 26:
2475
+ kv_cache_spec[f"layers.{i}.attn"] = MLAAttentionSpec(
2476
+ block_size=1024,
2477
+ num_kv_heads=1,
2478
+ head_size=576,
2479
+ dtype=torch.bfloat16,
2480
+ )
2481
+ else:
2482
+ kv_cache_spec[f"layers.{i}.linear_attn"] = new_mamba_spec(
2483
+ shapes=((512, 1024),), dtypes=(torch.float32,)
2484
+ )
2485
+ mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes
2486
+ assert kv_cache_spec["layers.0.linear_attn"].page_size_bytes > mla_page
2487
+
2488
+ with pytest.raises(NotImplementedError, match="page size"):
2489
+ kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec)
2490
+
2491
+
2492
+def test_get_kv_cache_config_mamba_hybrid_sharing_prepadded_mamba():
2493
+ """Platform-prepadded mamba pages (mamba_page_size_padded hint) must not
2494
+ disable slot sharing; the layout re-pads them to the MLA page."""
2495
+ model_config = ModelConfig(max_model_len=8192)
2496
+ vllm_config = VllmConfig(model_config=model_config)
2497
+
2498
+ def prepadded_mamba_spec():
2499
+ return new_mamba_spec(page_size_padded=294_912)
2500
+
2501
+ kv_cache_spec, _ = _glm5_like_kv_cache_spec(prepadded_mamba_spec)
2502
+ mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes
2503
+ assert prepadded_mamba_spec().page_size_bytes < mla_page
2504
+
2505
+ groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec)
2506
+ mamba_groups = [
2507
+ group for group in groups if isinstance(group.kv_cache_spec, MambaSpec)
2508
+ ]
2509
+ assert len(groups) == 5
2510
+ assert [len(group.layer_names) for group in mamba_groups] == [9, 9, 8, 8]
2511
+ for group in mamba_groups:
2512
+ assert group.kv_cache_spec.page_size_bytes == mla_page
2513
+
2514
+ bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups)
2515
+ kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups(
2516
+ vllm_config, groups, bytes_per_block * 100 + 1
2517
+ )
2518
+ assert kv_cache_config.num_blocks == 100
2519
+ assert len(kv_cache_config.kv_cache_tensors) == 56
2520
+
2521
+
2522
+def test_get_kv_cache_config_mamba_hybrid_sharing_pp_balanced_projection():
2523
+ """Round-robin mamba grouping keeps practical PP splits balanced: every
2524
+ stage's largest projected mamba group slice fits its projected MLA
2525
+ layers, so per-stage slot tensors all have an MLA owner."""
2526
+ model_config = ModelConfig(max_model_len=8192)
2527
+ vllm_config = VllmConfig(model_config=model_config)
2528
+
2529
+ kv_cache_spec, _ = _glm5_like_kv_cache_spec()
2530
+ global_groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec)
2531
+ mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes
2532
+ idx_page = kv_cache_spec["layers.3.indexer"].page_size_bytes
2533
+
2534
+ # PP=2 split at transformer layer 22/23.
2535
+ # Stage 0 of a PP=2 split at transformer layer 22/23: 5 MLA(+indexer)
2536
+ # layers, projected mamba groups [5, 5, 4, 4].
2537
+ worker_spec = {n: s for n, s in kv_cache_spec.items() if int(n.split(".")[1]) <= 22}
2538
+ groups = kv_cache_utils._project_kv_cache_groups_to_worker(
2539
+ global_groups, worker_spec
2540
+ )
2541
+ mamba_groups = [
2542
+ group for group in groups if isinstance(group.kv_cache_spec, MambaSpec)
2543
+ ]
2544
+ assert [len(group.layer_names) for group in mamba_groups] == [5, 5, 4, 4]
2545
+
2546
+ bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups)
2547
+ assert bytes_per_block == 5 * mla_page + 5 * idx_page
2548
+
2549
+ kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups(
2550
+ vllm_config, groups, bytes_per_block * 100 + 1
2551
+ )
2552
+ assert kv_cache_config.num_blocks == 100
2553
+ tensors = _tensor_by_layer(kv_cache_config)
2554
+ # Every projected Mamba slot aliases an MLA owner.
2555
+ assert tensors["layers.3.attn"].offset == tensors["layers.0.linear_attn"].offset
2556
+ assert tensors["layers.3.attn"].offset == tensors["layers.1.linear_attn"].offset
2557
+ assert tensors["layers.3.attn"].offset == tensors["layers.2.linear_attn"].offset
2558
+ assert tensors["layers.3.attn"].offset == tensors["layers.4.linear_attn"].offset
2559
+ # The last slot only hosts the two groups with a 5th projected layer.
2560
+ assert tensors["layers.19.attn"].offset == tensors["layers.21.linear_attn"].offset
2561
+ assert tensors["layers.19.attn"].offset == tensors["layers.22.linear_attn"].offset
2562
+ assert {t.size for t in kv_cache_config.kv_cache_tensors} == {bytes_per_block * 100}
2563
+
2564
+
2565
+def test_get_kv_cache_config_mamba_hybrid_sharing_pp_group_count_bump(monkeypatch):
2566
+ """PP=4's default partition [11,11,12,11] leaves stage 0 with 2 MLA
2567
+ layers but a round-robin slice of 3 under the minimum 4 mamba groups;
2568
+ grouping bumps to 5 groups so every stage's projection keeps sharing
2569
+ on instead of silently falling back on stage 0."""
2570
+ from vllm.config import ParallelConfig
2571
+ from vllm.distributed.utils import get_pp_indices
2572
+
2573
+ monkeypatch.setattr(ModelConfig, "get_total_num_hidden_layers", lambda self: 45)
2574
+ vllm_config = VllmConfig(
2575
+ model_config=ModelConfig(max_model_len=8192),
2576
+ parallel_config=ParallelConfig(pipeline_parallel_size=4),
2577
+ )
2578
+
2579
+ kv_cache_spec, mamba_layers = _glm5_like_kv_cache_spec()
2580
+ groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec)
2581
+ mamba_groups = [g for g in groups if isinstance(g.kv_cache_spec, MambaSpec)]
2582
+ assert [len(g.layer_names) for g in mamba_groups] == [7, 7, 7, 7, 6]
2583
+ for k, name in enumerate(mamba_layers):
2584
+ assert name in mamba_groups[k % 5].layer_names
2585
+
2586
+ for rank in range(4):
2587
+ start, end = get_pp_indices(45, rank, 4)
2588
+ worker_spec = {
2589
+ n: s
2590
+ for n, s in kv_cache_spec.items()
2591
+ if start <= int(n.split(".")[1]) < end
2592
+ }
2593
+ projected = kv_cache_utils._project_kv_cache_groups_to_worker(
2594
+ groups, worker_spec
2595
+ )
2596
+ assert kv_cache_utils._glm5_next_tensor_layout(projected) is not None
2597
+
2598
+
2599
+def test_get_kv_cache_config_mamba_hybrid_sharing_pp_starved_stage(monkeypatch):
2600
+ """Reject PP stages whose Mamba layers have no MLA slot to share."""
2601
+ from vllm.config import ParallelConfig
2602
+
2603
+ monkeypatch.setattr(ModelConfig, "get_total_num_hidden_layers", lambda self: 45)
2604
+ monkeypatch.setenv("VLLM_PP_LAYER_PARTITION", "3,42")
2605
+ vllm_config = VllmConfig(
2606
+ model_config=ModelConfig(max_model_len=8192),
2607
+ parallel_config=ParallelConfig(pipeline_parallel_size=2),
2608
)
2215
- scheduler_config = generate_scheduler_kv_cache_config([worker_config])
2609
2217
- assert worker_config.needs_kv_cache_zeroing
2218
- assert scheduler_config.needs_kv_cache_zeroing
2610
+ kv_cache_spec, _ = _glm5_like_kv_cache_spec()
2611
+ with pytest.raises(ValueError, match="VLLM_PP_LAYER_PARTITION"):
2612
+ kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec)
2613
+
2614
+
2615
+def test_get_kv_cache_config_mamba_hybrid_sharing_beats_cross_layers_flag():
2616
+ """Hybrid slot sharing must take precedence over the experimental
2617
+ enable_cross_layers_blocks packed layout: generic packing would give the
2618
+ MLA slots a strided view, breaking the contiguous virtual split."""
2619
+ model_config = ModelConfig(max_model_len=8192)
2620
+ kv_transfer_config = KVTransferConfig(
2621
+ kv_connector="NixlConnector",
2622
+ kv_role="kv_both",
2623
+ kv_connector_extra_config={"enable_cross_layers_blocks": "true"},
2624
+ )
2625
+ vllm_config = VllmConfig(
2626
+ model_config=model_config, kv_transfer_config=kv_transfer_config
2627
+ )
2628
+
2629
+ kv_cache_spec, _ = _glm5_like_kv_cache_spec()
2630
+ groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec)
2631
+
2632
+ mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes
2633
+ idx_page = kv_cache_spec["layers.3.indexer"].page_size_bytes
2634
+ bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups)
2635
+ assert bytes_per_block == 11 * mla_page + 11 * idx_page
2636
+
2637
+ kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups(
2638
+ vllm_config, groups, bytes_per_block * 100 + 1
2639
+ )
2640
+ assert kv_cache_config.num_blocks == 100
2641
+ assert len(kv_cache_config.kv_cache_tensors) == 56
2642
+ assert all(
2643
+ tensor.block_stride in (mla_page, idx_page)
2644
+ for tensor in kv_cache_config.kv_cache_tensors
2645
+ )
2646
+ assert {t.size for t in kv_cache_config.kv_cache_tensors} == {bytes_per_block * 100}
2647
+
2648
+
2649
+def test_get_kv_cache_config_mamba_hybrid_sharing_no_indexer():
2650
+ """Kimi-Linear-like: MLA without indexer layers, idx_stride == 0."""
2651
+ model_config = ModelConfig(max_model_len=8192)
2652
+ vllm_config = VllmConfig(model_config=model_config)
2653
+ vllm_config.cache_config.kv_cache_layout = "LBNHC"
2654
+
2655
+ # 20 mamba + 7 MLA layers -> G = ceil(20 / 7) = 3 groups of [7, 7, 6].
2656
+ kv_cache_spec: dict[str, KVCacheSpec] = {}
2657
+ for i in range(27):
2658
+ if i % 4 == 3 or i == 26:
2659
+ kv_cache_spec[f"layers.{i}.attn"] = MLAAttentionSpec(
2660
+ block_size=1024,
2661
+ num_kv_heads=1,
2662
+ head_size=576,
2663
+ dtype=torch.bfloat16,
2664
+ )
2665
+ else:
2666
+ kv_cache_spec[f"layers.{i}.linear_attn"] = new_mamba_spec()
2667
+ mla_page = kv_cache_spec["layers.3.attn"].page_size_bytes
2668
+
2669
+ groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec)
2670
+ mamba_groups = [
2671
+ group for group in groups if isinstance(group.kv_cache_spec, MambaSpec)
2672
+ ]
2673
+ assert len(groups) == 4
2674
+ assert [len(group.layer_names) for group in mamba_groups] == [7, 7, 6]
2675
+ for group in mamba_groups:
2676
+ assert group.kv_cache_spec.page_size_bytes == mla_page
2677
+
2678
+ bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups)
2679
+ assert bytes_per_block == 7 * mla_page
2680
+
2681
+ kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups(
2682
+ vllm_config, groups, bytes_per_block * 100 + 1
2683
+ )
2684
+ assert kv_cache_config.num_blocks == 100
2685
+ # Each group has one tensor descriptor and all groups alias the same slots.
2686
+ assert len(kv_cache_config.kv_cache_tensors) == 4
2687
+ tensors = _tensor_by_layer(kv_cache_config)
2688
+ mla_names = [f"layers.{i}.attn" for i in (*range(3, 27, 4), 26)]
2689
+ for index, mla_name in enumerate(mla_names):
2690
+ mla_offset = _layer_offset(tensors[mla_name], mla_name)
2691
+ for group in mamba_groups:
2692
+ if index < len(group.layer_names):
2693
+ mamba_name = group.layer_names[index]
2694
+ assert _layer_offset(tensors[mamba_name], mamba_name) == mla_offset
2695
+ assert {t.size for t in kv_cache_config.kv_cache_tensors} == {bytes_per_block * 100}
2696
+
2697
+
2698
+def test_get_kv_cache_capacity_after_scheduler_unwrap():
2699
+ """max_concurrency must survive the scheduler-config unwrap.
2700
+
2701
+ Regression for the balanced-mamba hybrid layout:
2702
+ ``generate_scheduler_kv_cache_config`` flattens the MLA
2703
+ ``UniformTypeKVCacheSpecs`` group into a single ``MLAAttentionSpec``, so the
2704
+ scheduler config holds an MLA spec (page size A) next to ``MambaSpec``
2705
+ groups (page size B). ``EngineCore._initialize_kv_caches`` calls
2706
+ ``get_kv_cache_capacity`` on that unwrapped config; the uniform-page-size
2707
+ path used to assert-fail (``assert len(page_sizes) == 1``) on this topology.
2708
+ """
2709
+ model_config = ModelConfig(max_model_len=8192)
2710
+ vllm_config = VllmConfig(model_config=model_config)
2711
+
2712
+ kv_cache_spec: dict[str, KVCacheSpec] = {}
2713
+ for i in range(45):
2714
+ if i % 4 == 3:
2715
+ kv_cache_spec[f"layers.{i}.attn"] = MLAAttentionSpec(
2716
+ block_size=1024,
2717
+ num_kv_heads=1,
2718
+ head_size=576,
2719
+ dtype=torch.bfloat16,
2720
+ )
2721
+ kv_cache_spec[f"layers.{i}.indexer"] = MLAAttentionSpec(
2722
+ block_size=1024,
2723
+ num_kv_heads=1,
2724
+ head_size=132,
2725
+ dtype=torch.uint8,
2726
+ tokens_per_state=16,
2727
+ )
2728
+ else:
2729
+ kv_cache_spec[f"layers.{i}.linear_attn"] = new_mamba_spec()
2730
+
2731
+ groups = kv_cache_utils.get_kv_cache_groups(vllm_config, kv_cache_spec)
2732
+ bytes_per_block = kv_cache_utils._pool_bytes_per_block(groups)
2733
+ kv_cache_config = kv_cache_utils.get_kv_cache_config_from_groups(
2734
+ vllm_config, groups, bytes_per_block * 100 + 1
2735
+ )
2736
+
2737
+ scheduler_config = generate_scheduler_kv_cache_config([kv_cache_config])
2738
+ # Confirm we are exercising the regression path: the MLA group is no longer
2739
+ # UniformTypeKVCacheSpecs, so sibling specs have differing page sizes.
2740
+ assert not any(
2741
+ isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs)
2742
+ for g in scheduler_config.kv_cache_groups
2743
+ )
2744
+
2745
+ unwrapped_groups = scheduler_config.kv_cache_groups
2746
+ expected_max_mem = kv_cache_utils._max_memory_usage_bytes_from_groups(
2747
+ vllm_config, unwrapped_groups
2748
+ )
2749
+ expected_pool = kv_cache_utils._pool_bytes_per_block(unwrapped_groups)
2750
+ expected_blocks_per_request = (
2751
+ expected_max_mem + expected_pool - 1
2752
+ ) // expected_pool
2753
+
2754
+ _, max_concurrency = get_kv_cache_capacity(vllm_config, scheduler_config)
2755
+ assert max_concurrency > 0
2756
+ assert max_concurrency == pytest.approx(
2757
+ scheduler_config.num_blocks / expected_blocks_per_request
2758
+ )
2759
2760
2221
-def new_mla_spec(cache_dtype_str=None, block_size=16):
2761
+def new_mla_spec(cache_dtype_str=None, block_size: int = 16):
2762
# head_size = kv_lora_rank(512) + qk_rope_head_dim(64) = 576
2763
return MLAAttentionSpec(
2764
block_size=block_size,
@@ -3305,9 +3845,10 @@ def _hybrid_specs_with_draft(draft: bool, draft_shares_target_spec: bool = False
3845
"""A K3-shaped hybrid: MLA full attention + Mamba, optionally plus a
3846
DSpark-style draft MLA layer marked non_causal_multi_token_decode.
3847
3308
- The target's fp8 KV dtype is what keeps the draft in its own bucket, as it
3309
- does on Kimi-K3 (target `--kv-cache-dtype fp8_e4m3`, draft `auto`). Pass
3310
- draft_shares_target_spec to collapse them into one group instead.
3848
+ The target's fp8 KV dtype is one difference from the draft on Kimi-K3
3849
+ (target `--kv-cache-dtype fp8_e4m3`, draft `auto`). Pass
3850
+ draft_shares_target_spec to match that dtype while retaining the draft
3851
+ marker as a group boundary.
3852
"""
3853
target_dtype = None if draft_shares_target_spec else "fp8_e4m3"
3854
specs = {
tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py
+193
@@ -30,6 +30,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.worker import (
30
from vllm.v1.core.kv_cache_utils import BlockHash
31
from vllm.v1.kv_cache_interface import (
32
FullAttentionSpec,
33
+ KpoolTailSpec,
34
KVCacheConfig,
35
KVCacheGroupSpec,
36
KVCacheTensor,
@@ -605,3 +606,195 @@ def test_sub_block_partial_tail_offload_covers_smaller_group_blocks():
606
token_dbs[1].key_for(hs[2]): [10_000 + mamba_cow_block * 512],
607
}
608
assert store.puts == expected
609
+
610
+
611
+def test_worker_lookup_hits_sub_block_partial_tail():
612
+ """worker.lookup must query sub-block keys when partial hash hits are on.
613
+
614
+ Regression test: ``lookup`` hard-coded ``fine_grained = False``, so the
615
+ sub-block keys persisted by ``_sub_block_tail_puts`` were never probed and
616
+ partial prefix hits silently returned 0 while the store side kept writing
617
+ them. Here the mamba block (16) exceeds the hash unit (4), so
618
+ ``enable_partial_hash_hits`` is on and the lookup must find the stored
619
+ boundary at 12.
620
+ """
621
+ full = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None)
622
+ mamba = MambaSpec(
623
+ block_size=16,
624
+ shapes=((1, 1),),
625
+ dtypes=(torch.float32,),
626
+ mamba_cache_mode="align",
627
+ )
628
+ cfg = KVCacheConfig(
629
+ num_blocks=4,
630
+ kv_cache_tensors=[
631
+ KVCacheTensor(
632
+ size=4 * full.page_size_bytes,
633
+ layers=["L0"],
634
+ layer_stride=4 * full.page_size_bytes,
635
+ block_stride=full.page_size_bytes,
636
+ ),
637
+ KVCacheTensor(
638
+ size=4 * mamba.page_size_bytes,
639
+ layers=["L1"],
640
+ layer_stride=4 * mamba.page_size_bytes,
641
+ block_stride=mamba.page_size_bytes,
642
+ ),
643
+ ],
644
+ kv_cache_groups=[
645
+ KVCacheGroupSpec(["L0"], full),
646
+ KVCacheGroupSpec(["L1"], mamba),
647
+ ],
648
+ )
649
+ vllm_config = _minimal_vllm_config(cache_block_size=16)
650
+ # Hash unit 4 < mamba block 16 -> partial hash hits are enabled.
651
+ vllm_config.cache_config.prefix_match_unit = 4
652
+ store = _DictStore()
653
+
654
+ worker = _build_worker_with_dict_store(vllm_config, cfg, store)
655
+ worker.tp_size = 1
656
+ worker.pp_size = 1
657
+ worker.num_kv_head = 8
658
+ assert worker.coord.enable_partial_hash_hits
659
+
660
+ for g_idx, db in enumerate(worker.token_dbs):
661
+ db.set_kv_caches_base_addr([g_idx * 10_000])
662
+ db.set_block_len([512])
663
+
664
+ send_thread = KVCacheStoreSendingThread(
665
+ store=store,
666
+ token_databases=worker.token_dbs,
667
+ block_size=worker.block_size,
668
+ coord=worker.coord,
669
+ tp_rank=0,
670
+ group_put_steps=[1, 1],
671
+ kv_role="kv_both",
672
+ ready_event=threading.Event(),
673
+ replicate_config=MagicMock(),
674
+ )
675
+
676
+ # Persist the sub-block partial tail at boundary 12 (keyed by hs[12//4-1]).
677
+ hs = [BlockHash(bytes([i + 1]) * 4) for i in range(5)]
678
+ req = ReqMeta(
679
+ req_id="r0",
680
+ token_len_chunk=0,
681
+ block_ids=([1], [2]),
682
+ block_hashes=hs,
683
+ can_save=True,
684
+ num_prompt_tokens=20,
685
+ boundary_state_offloads=[(1, 7, 12)],
686
+ )
687
+ send_thread._maybe_offload_boundary_states(req)
688
+
689
+ worker.store = store
690
+
691
+ # A 13-token prompt sharing the prefix must hit the stored boundary at 12.
692
+ assert worker.lookup(num_tokens=13, block_hashes=hs).hit_length == 12
693
+
694
+
695
+def test_worker_setup_tolerates_finer_scratch_group():
696
+ """Setup and lookup must tolerate a non-prefix-cacheable scratch group.
697
+
698
+ Regression: GLM-5.3-Flash carries a kpool-tail scratch group whose block
699
+ size (``index_kpool`` tokens) is finer than the hash unit, so the
700
+ coordinator's divisibility assert and the scratch group's
701
+ ``ChunkedTokenDatabase`` both rejected worker setup for any hash unit the
702
+ scratch block does not divide. Scratch groups never participate in
703
+ store/load/lookup, so setup must skip them and lookup must still hit
704
+ stored sub-block boundaries on the participating groups.
705
+ """
706
+ full = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None)
707
+ mamba = MambaSpec(
708
+ block_size=16,
709
+ shapes=((1, 1),),
710
+ dtypes=(torch.float32,),
711
+ mamba_cache_mode="align",
712
+ )
713
+ # Scratch block 4 is not divisible by the hash unit 8 below.
714
+ scratch = KpoolTailSpec(
715
+ block_size=4,
716
+ num_kv_heads=2,
717
+ head_size=64,
718
+ head_size_v=0,
719
+ dtype=torch.bfloat16,
720
+ sliding_window=4,
721
+ )
722
+ cfg = KVCacheConfig(
723
+ num_blocks=4,
724
+ kv_cache_tensors=[
725
+ KVCacheTensor(
726
+ size=4 * full.page_size_bytes,
727
+ layers=["L0"],
728
+ layer_stride=4 * full.page_size_bytes,
729
+ block_stride=full.page_size_bytes,
730
+ ),
731
+ KVCacheTensor(
732
+ size=4 * mamba.page_size_bytes,
733
+ layers=["L1"],
734
+ layer_stride=4 * mamba.page_size_bytes,
735
+ block_stride=mamba.page_size_bytes,
736
+ ),
737
+ KVCacheTensor(
738
+ size=4 * scratch.page_size_bytes,
739
+ layers=["L2"],
740
+ layer_stride=4 * scratch.page_size_bytes,
741
+ block_stride=scratch.page_size_bytes,
742
+ ),
743
+ ],
744
+ kv_cache_groups=[
745
+ KVCacheGroupSpec(["L0"], full),
746
+ KVCacheGroupSpec(["L1"], mamba),
747
+ KVCacheGroupSpec(["L2"], scratch),
748
+ ],
749
+ )
750
+ vllm_config = _minimal_vllm_config(cache_block_size=16)
751
+ # Hash unit 8 divides the participating groups (16) but not the scratch
752
+ # group (4); mamba block 16 > 8 keeps partial hash hits on.
753
+ vllm_config.cache_config.prefix_match_unit = 8
754
+ store = _DictStore()
755
+
756
+ worker = _build_worker_with_dict_store(vllm_config, cfg, store)
757
+ worker.tp_size = 1
758
+ worker.pp_size = 1
759
+ worker.num_kv_head = 8
760
+ assert worker.coord.enable_partial_hash_hits
761
+ # The scratch DB is keyed at its own block size and never probed.
762
+ assert worker.token_dbs[2].hash_block_size == 4
763
+
764
+ for g_idx, db in enumerate(worker.token_dbs):
765
+ db.set_kv_caches_base_addr([g_idx * 10_000])
766
+ db.set_block_len([512])
767
+
768
+ send_thread = KVCacheStoreSendingThread(
769
+ store=store,
770
+ token_databases=worker.token_dbs,
771
+ block_size=worker.block_size,
772
+ coord=worker.coord,
773
+ tp_rank=0,
774
+ group_put_steps=[1, 1, 1],
775
+ kv_role="kv_both",
776
+ ready_event=threading.Event(),
777
+ replicate_config=MagicMock(),
778
+ group_participates=[True, True, False],
779
+ )
780
+
781
+ # Persist the sub-block partial tail at boundary 12 (keyed by hs[12//8-1]).
782
+ hs = [BlockHash(bytes([i + 1]) * 8) for i in range(3)]
783
+ req = ReqMeta(
784
+ req_id="r0",
785
+ token_len_chunk=0,
786
+ block_ids=([1], [2], [3]),
787
+ block_hashes=hs,
788
+ can_save=True,
789
+ num_prompt_tokens=20,
790
+ boundary_state_offloads=[(1, 7, 12)],
791
+ )
792
+ send_thread._maybe_offload_boundary_states(req)
793
+
794
+ worker.store = store
795
+
796
+ # A 13-token prompt sharing the prefix must hit the first hash unit.
797
+ assert worker.lookup(num_tokens=13, block_hashes=hs).hit_length == 8
798
+ # The scratch group's namespace never enters the store.
799
+ scratch_prefix = worker.token_dbs[2].key_for(hs[0]).rsplit("@", 1)[0]
800
+ assert not any(key.startswith(scratch_prefix) for key in store._data)
tests/v1/kv_connector/unit/test_nixl_desc_geometry.py
+161
@@ -211,6 +211,167 @@ def _make_mla_hybrid_worker(local_block_size, kernel_block_size, num_logical_blo
211
return worker
212
213
214
+@pytest.mark.cpu_test
215
+@pytest.mark.parametrize("logical_block_size", [1152, 640])
216
+@pytest.mark.parametrize("tail_first", [False, True])
217
+def test_register_compressed_indexer_uses_virtual_transfer_pages(
218
+ logical_block_size, tail_first
219
+):
220
+ """Compressed indexer rows must split into contiguous NIXL transfer pages."""
221
+ from unittest.mock import MagicMock
222
+
223
+ from vllm.config import set_current_vllm_config
224
+ from vllm.distributed.kv_transfer.kv_connector.v1.nixl import (
225
+ base_worker as bw,
226
+ )
227
+ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
228
+ NixlConnectorWorker,
229
+ )
230
+ from vllm.v1.kv_cache_interface import (
231
+ KpoolTailSpec,
232
+ KVCacheConfig,
233
+ KVCacheGroupSpec,
234
+ KVCacheLayout,
235
+ KVCacheTensor,
236
+ MLAAttentionSpec,
237
+ UniformTypeKVCacheSpecs,
238
+ create_kv_cache_views,
239
+ )
240
+
241
+ num_logical_blocks = 3
242
+ transfer_block_size = 64
243
+ kernel_block_size = 128
244
+ tokens_per_state = 4
245
+ state_content_bytes = 132
246
+
247
+ indexer_spec = MLAAttentionSpec(
248
+ block_size=logical_block_size,
249
+ num_kv_heads=1,
250
+ head_size=128,
251
+ head_size_v=0,
252
+ dtype=torch.uint8,
253
+ state_content_bytes=state_content_bytes,
254
+ tokens_per_state=tokens_per_state,
255
+ )
256
+ indexer_page_size = indexer_spec.page_size_bytes
257
+ tail_spec = KpoolTailSpec(
258
+ block_size=tokens_per_state,
259
+ num_kv_heads=2,
260
+ head_size=128,
261
+ head_size_v=0,
262
+ dtype=torch.bfloat16,
263
+ page_size_padded=indexer_page_size,
264
+ sliding_window=tokens_per_state,
265
+ )
266
+
267
+ allocation_size = num_logical_blocks * indexer_page_size
268
+ indexer_tensor = KVCacheTensor(
269
+ size=allocation_size,
270
+ layers=["indexer"],
271
+ layer_stride=allocation_size,
272
+ block_stride=indexer_page_size,
273
+ )
274
+ tail_tensor = KVCacheTensor(
275
+ size=allocation_size,
276
+ layers=["tail"],
277
+ layer_stride=allocation_size,
278
+ block_stride=indexer_page_size,
279
+ )
280
+ kv_cache_config = KVCacheConfig(
281
+ num_blocks=num_logical_blocks,
282
+ kv_cache_tensors=[indexer_tensor, tail_tensor],
283
+ kv_cache_groups=[
284
+ KVCacheGroupSpec(
285
+ ["indexer"],
286
+ UniformTypeKVCacheSpecs(
287
+ block_size=logical_block_size,
288
+ kv_cache_specs={"indexer": indexer_spec},
289
+ ),
290
+ ),
291
+ KVCacheGroupSpec(
292
+ ["tail"],
293
+ UniformTypeKVCacheSpecs(
294
+ block_size=tokens_per_state,
295
+ kv_cache_specs={"tail": tail_spec},
296
+ ),
297
+ ),
298
+ ],
299
+ )
300
+
301
+ raw = torch.zeros(allocation_size, dtype=torch.int8)
302
+ (indexer_cache,) = create_kv_cache_views(
303
+ raw,
304
+ indexer_spec,
305
+ num_logical_blocks,
306
+ KVCacheLayout.LBHNC,
307
+ indexer_tensor,
308
+ kernel_block_size=kernel_block_size,
309
+ )
310
+ (tail_cache,) = create_kv_cache_views(
311
+ raw,
312
+ tail_spec,
313
+ num_logical_blocks,
314
+ KVCacheLayout.LBHNC,
315
+ tail_tensor,
316
+ )
317
+ assert indexer_cache.data_ptr() == tail_cache.data_ptr() == raw.data_ptr()
318
+
319
+ vllm_config = create_vllm_config(block_size=logical_block_size)
320
+ vllm_config.cache_config.kv_cache_layout = "LBHNC"
321
+ vllm_config.kv_transfer_config.kv_buffer_device = "cuda"
322
+ fake_backend = MagicMock()
323
+ fake_backend.get_supported_kernel_block_sizes.return_value = [transfer_block_size]
324
+ fake_backend.get_name.return_value = "DEEPSEEK_V32_INDEXER"
325
+ fake_backend.full_cls_name.return_value = "fake.DEEPSEEK_V32_INDEXER"
326
+ fake_platform = MagicMock()
327
+ fake_platform.device_type = "cuda"
328
+ fake_platform.get_nixl_memory_type.return_value = "VRAM"
329
+
330
+ caches = [("indexer", indexer_cache), ("tail", tail_cache)]
331
+ if tail_first:
332
+ caches.reverse()
333
+
334
+ with (
335
+ patch.object(bw, "NixlWrapper", _RecordingNixl),
336
+ patch.object(bw, "get_tensor_model_parallel_rank", return_value=0),
337
+ patch.object(bw, "get_tensor_model_parallel_world_size", return_value=1),
338
+ patch.object(bw, "get_current_attn_backends", return_value=[fake_backend]),
339
+ patch.object(bw, "current_platform", fake_platform),
340
+ set_current_vllm_config(vllm_config),
341
+ ):
342
+ worker = NixlConnectorWorker(vllm_config, "local-engine", kv_cache_config)
343
+ worker.use_mla = True
344
+ worker.register_kv_caches(dict(caches))
345
+
346
+ transfer_page_size = transfer_block_size // tokens_per_state * state_content_bytes
347
+ num_transfer_blocks = num_logical_blocks * (
348
+ logical_block_size // transfer_block_size
349
+ )
350
+ expected_descs = np.asarray(
351
+ [
352
+ [
353
+ raw.data_ptr() + block_idx * transfer_page_size,
354
+ transfer_page_size,
355
+ 0,
356
+ ]
357
+ for block_idx in range(num_transfer_blocks)
358
+ ],
359
+ dtype=np.uint64,
360
+ )
361
+
362
+ assert worker.block_size == transfer_block_size
363
+ assert worker.num_regions == 1
364
+ assert worker.block_len_per_layer == [transfer_page_size]
365
+ assert worker.block_stride_per_layer == [transfer_page_size]
366
+ assert worker._region_is_mla == [True]
367
+ assert worker.kv_caches_base_addr[worker.engine_id][0] == [raw.data_ptr()]
368
+ assert worker._registered_descs[0] == [(raw.data_ptr(), raw.nbytes, 0, "")]
369
+ np.testing.assert_array_equal(worker.src_blocks_data, expected_descs)
370
+ assert expected_descs[-1, 0] + expected_descs[-1, 1] == (
371
+ raw.data_ptr() + raw.nbytes
372
+ )
373
+
374
+
375
def _make_remote_meta(
376
worker,
377
remote_block_size,
tests/v1/worker/test_attn_utils.py
+12
-4
@@ -143,15 +143,22 @@ def test_reshape_padded_kv_cache_strides_by_padded_page():
143
144
145
@pytest.mark.parametrize(
146
- ("kernel_block_sizes", "expected_num_blocks", "expected_num_states"),
146
+ (
147
+ "kernel_block_sizes",
148
+ "storage_block_size",
149
+ "expected_num_blocks",
150
+ "expected_num_states",
151
+ ),
152
[
148
- (None, 4, 64),
149
- ([256], 4, 64),
150
- ([64], 16, 16),
153
+ (None, None, 4, 64),
154
+ ([256], None, 4, 64),
155
+ ([64], None, 16, 16),
156
+ ([64], 256, 4, 64),
157
],
158
)
159
def test_allocate_compressed_mla_cache(
160
kernel_block_sizes: list[int] | None,
161
+ storage_block_size: int | None,
162
expected_num_blocks: int,
163
expected_num_states: int,
164
):
@@ -161,6 +168,7 @@ def test_allocate_compressed_mla_cache(
168
head_size=128,
169
dtype=torch.bfloat16,
170
tokens_per_state=4,
171
+ storage_block_size=storage_block_size,
172
)
173
num_pages = 4
174
config = KVCacheConfig(
tests/v1/worker/test_dsv4_packed_zeroer_geometry.py
+2
@@ -76,6 +76,7 @@ def test_packed_dsv4_zeroer_zeroes_only_each_layers_page():
76
static_forward_context={
77
f"layer.{i}": SimpleNamespace(kv_cache=views[i]) for i in range(NUM_LAYERS)
78
},
79
+ num_blocks=NUM_BLOCKS,
80
)
81
seg_addrs, seg_block_strides, seg_page_sizes, _, _, n_segs = zeroer._meta
82
@@ -153,6 +154,7 @@ def test_overlaid_zeroer_dedups_segments_with_max_span():
154
static_forward_context={
155
name: SimpleNamespace(kv_cache=views[name]) for name in views
156
},
157
+ num_blocks=config.num_blocks,
158
)
159
seg_addrs, seg_block_strides, seg_page_sizes, _, _, n_segs = zeroer._meta
160
tests/v1/worker/test_kv_block_zeroer.py
+53
@@ -54,6 +54,7 @@ def test_attention_blocks_are_zeroed(spec):
54
static_forward_context={
55
layer_name: SimpleNamespace(kv_cache=storage),
56
},
57
+ num_blocks=4,
58
)
59
60
zeroer.zero_block_ids([1])
@@ -64,6 +65,56 @@ def test_attention_blocks_are_zeroed(spec):
65
assert torch.equal(storage, expected)
66
67
68
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
69
+def test_layers_in_one_group_may_use_different_kernel_pages_per_block():
70
+ """Derive kernel pages per block from each layer's allocation."""
71
+ device = torch.device("cuda")
72
+ num_blocks = 4
73
+ spec = SlidingWindowSpec(
74
+ block_size=2,
75
+ num_kv_heads=1,
76
+ head_size=1,
77
+ dtype=torch.int32,
78
+ sliding_window=2,
79
+ )
80
+ # Two kernel pages per logical block.
81
+ wide = torch.ones((num_blocks * 2, 3), dtype=torch.int32, device=device)
82
+ # One kernel page per logical block, carved out of a larger allocation so
83
+ # an over-strided write lands in the guard region instead of faulting or
84
+ # silently hitting another tensor.
85
+ narrow_backing = torch.ones((num_blocks * 3, 5), dtype=torch.int32, device=device)
86
+ narrow = narrow_backing[:num_blocks]
87
+
88
+ zeroer = KVBlockZeroer(
89
+ device,
90
+ attn_groups_iter=[
91
+ AttentionGroup(
92
+ None,
93
+ ["wide", "narrow"],
94
+ spec,
95
+ 0,
96
+ )
97
+ ],
98
+ kernel_block_sizes=[1],
99
+ static_forward_context={
100
+ "wide": SimpleNamespace(kv_cache=wide),
101
+ "narrow": SimpleNamespace(kv_cache=narrow),
102
+ },
103
+ num_blocks=num_blocks,
104
+ )
105
+
106
+ zeroer.zero_block_ids([num_blocks - 1])
107
+ torch.accelerator.synchronize()
108
+
109
+ expected_wide = torch.ones_like(wide)
110
+ expected_wide[2 * (num_blocks - 1) :] = 0
111
+ assert torch.equal(wide, expected_wide)
112
+
113
+ expected_narrow = torch.ones_like(narrow_backing)
114
+ expected_narrow[num_blocks - 1] = 0
115
+ assert torch.equal(narrow_backing, expected_narrow)
116
+
117
+
118
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
119
def test_block_ids_are_not_overwritten_while_copy_is_in_flight():
120
device = torch.device("cuda")
@@ -207,6 +258,7 @@ def test_large_dsv4_launch_geometry(monkeypatch):
258
name: SimpleNamespace(kv_cache=storage)
259
for name, storage in storages.items()
260
},
261
+ num_blocks=1,
262
)
263
264
assert zeroer._meta is not None
@@ -332,6 +384,7 @@ def test_zeroes_exactly_one_block_per_layer(layout: KVCacheLayout):
384
attn_groups_iter=iter(groups),
385
kernel_block_sizes=[spec.block_size],
386
static_forward_context=ctx,
387
+ num_blocks=num_blocks,
388
)
389
zeroer.zero_block_ids([2])
390
torch.accelerator.synchronize()
tests/v1/worker/test_mamba_hybrid_model_state.py
+42
@@ -7,15 +7,57 @@ from unittest.mock import Mock
7
import pytest
8
import torch
9
10
+from vllm.config.compilation import CUDAGraphMode
11
from vllm.platforms import current_platform
12
from vllm.v1.attention.backends.recoverssm_metadata import (
13
RecoverSSMMetadata,
14
RecoverSSMPostprocessMetadata,
15
)
16
+from vllm.v1.worker.gpu.model_states import mamba_hybrid
17
from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState
18
from vllm.v1.worker.gpu.model_states.recoverssm import RecoverSSMState
19
20
21
+def test_prepare_attn_forwards_positions(monkeypatch: pytest.MonkeyPatch) -> None:
22
+ state = object.__new__(MambaHybridModelState)
23
+ state.vllm_config = SimpleNamespace(num_speculative_tokens=0)
24
+ state.max_model_len = 8192
25
+ state._align_mode = False
26
+ state.recoverssm = None
27
+
28
+ positions = torch.tensor([1536], dtype=torch.int64)
29
+ input_batch = SimpleNamespace(
30
+ num_reqs=1,
31
+ num_tokens=1,
32
+ num_reqs_after_padding=1,
33
+ num_tokens_after_padding=1,
34
+ query_start_loc_np=torch.tensor([0, 1], dtype=torch.int32).numpy(),
35
+ query_start_loc=torch.tensor([0, 1], dtype=torch.int32),
36
+ num_scheduled_tokens=torch.tensor([1], dtype=torch.int32),
37
+ seq_lens_cpu_upper_bound=torch.tensor([1537], dtype=torch.int32),
38
+ seq_lens=torch.tensor([1537], dtype=torch.int32),
39
+ is_prefilling_np=torch.tensor([False]).numpy(),
40
+ dcp_local_seq_lens=None,
41
+ positions=positions,
42
+ prompt_lens=torch.tensor([1024], dtype=torch.int32),
43
+ )
44
+ expected_metadata = {"layer": object()}
45
+ build_attn_metadata = Mock(return_value=expected_metadata)
46
+ monkeypatch.setattr(mamba_hybrid, "build_attn_metadata", build_attn_metadata)
47
+
48
+ metadata = state.prepare_attn(
49
+ input_batch=input_batch,
50
+ cudagraph_mode=CUDAGraphMode.NONE,
51
+ block_tables=(),
52
+ slot_mappings=torch.empty(0, dtype=torch.int64),
53
+ attn_groups=[],
54
+ kv_cache_config=Mock(),
55
+ )
56
+
57
+ assert metadata is expected_metadata
58
+ assert build_attn_metadata.call_args.kwargs["positions"] is positions
59
+
60
+
61
@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA")
62
@pytest.mark.parametrize(("num_sampled", "expected_value"), [(0, 1), (3, 3)])
63
def test_postprocess_state_scalar_with_int32_mapping(
vllm/_aiter_ops.py
+24
-19
@@ -3468,7 +3468,7 @@ class rocm_aiter_ops:
3468
# AITER's Python wrapper allocates intermediate/output tensors without
3469
# explicit device arguments, so run it under the residual tensor's device.
3470
with torch.device(residual_flat.device):
3471
- post_mix, comb_mix, layer_input = mhc_pre(
3471
+ args = (
3472
residual_flat,
3473
fn,
3474
hc_scale,
@@ -3478,9 +3478,11 @@ class rocm_aiter_ops:
3478
hc_sinkhorn_eps,
3479
hc_post_mult_value,
3480
sinkhorn_repeat,
3481
- norm_weight,
3482
- norm_eps,
3481
)
3482
+ if norm_weight is None:
3483
+ post_mix, comb_mix, layer_input = mhc_pre(*args)
3484
+ else:
3485
+ post_mix, comb_mix, layer_input = mhc_pre(*args, norm_weight, norm_eps)
3486
return (
3487
post_mix.view(*outer_shape, hc_mult, 1),
3488
comb_mix.view(*outer_shape, hc_mult, hc_mult),
@@ -3630,23 +3632,26 @@ class rocm_aiter_ops:
3632
),
3633
)
3634
3635
+ args = (
3636
+ x_flat,
3637
+ residual_flat,
3638
+ post_flat,
3639
+ comb_flat,
3640
+ fn,
3641
+ hc_scale,
3642
+ hc_base,
3643
+ rms_eps,
3644
+ hc_pre_eps,
3645
+ hc_sinkhorn_eps,
3646
+ hc_post_mult_value,
3647
+ sinkhorn_repeat,
3648
+ )
3649
with torch.device(residual_flat.device):
3634
- post_mix, comb_mix, layer_input, next_residual = mhc_fused_post_pre(
3635
- x_flat,
3636
- residual_flat,
3637
- post_flat,
3638
- comb_flat,
3639
- fn,
3640
- hc_scale,
3641
- hc_base,
3642
- rms_eps,
3643
- hc_pre_eps,
3644
- hc_sinkhorn_eps,
3645
- hc_post_mult_value,
3646
- sinkhorn_repeat,
3647
- norm_weight,
3648
- norm_eps,
3649
- )
3650
+ if norm_weight is None:
3651
+ result = mhc_fused_post_pre(*args)
3652
+ else:
3653
+ result = mhc_fused_post_pre(*args, norm_weight, norm_eps)
3654
+ post_mix, comb_mix, layer_input, next_residual = result
3655
3656
return (
3657
next_residual.view_as(residual),
vllm/config/speculative.py
+7
@@ -61,6 +61,7 @@ MTPModelTypes = Literal[
61
"hy_v4_mtp",
62
"gemma4_mtp",
63
"inkling_mtp",
64
+ "glm5_next_mtp",
65
]
66
NgramGPUTypes = Literal["ngram_gpu"]
67
DFlashModelTypes = Literal["dflash"]
@@ -1016,6 +1017,12 @@ class SpeculativeConfig:
1017
hf_config.update(
1018
{"n_predict": n_predict, "architectures": ["MiniMaxM3MTP"]}
1019
)
1020
+ if hf_config.model_type == "glm5_next":
1021
+ hf_config.model_type = "glm5_next_mtp"
1022
+ n_predict = hf_config.num_nextn_predict_layers
1023
+ hf_config.update(
1024
+ {"n_predict": n_predict, "architectures": ["Glm5NextMTPModel"]}
1025
+ )
1026
1027
return hf_config
1028
vllm/config/vllm.py
+3
@@ -81,6 +81,9 @@ DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES = frozenset(
81
"DeepSeekV4MTPModel",
82
"Dots3NoteForCausalLM",
83
"Dots3NoteMTPModel",
84
+ "Glm5NextForCausalLM",
85
+ "Glm5NextForConditionalGeneration",
86
+ "Glm5NextMTPModel",
87
"GlmMoeDsaForCausalLM",
88
"HYV4ForCausalLM",
89
"HYV4MTPModel",
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py
+4
-1
@@ -56,6 +56,7 @@ from vllm.v1.core.sched.output import SchedulerOutput
56
from vllm.v1.kv_cache_interface import (
57
AttentionSpec,
58
FullAttentionSpec,
59
+ KpoolTailSpec,
60
KVCacheSpec,
61
MambaSpec,
62
SlidingWindowSpec,
@@ -1675,7 +1676,9 @@ class MooncakeConnectorWorker:
1676
block_len = region_cache.stride(0) * region_cache.element_size()
1677
region_base_addresses.append(base_addr)
1678
1678
- if isinstance(layer_spec, AttentionSpec) and block_is_contiguous:
1679
+ if isinstance(layer_spec, KpoolTailSpec):
1680
+ kv_block_len = layer_spec.unpadded_page_size_bytes // 2
1681
+ elif isinstance(layer_spec, AttentionSpec) and block_is_contiguous:
1682
assert (
1683
layer_spec.page_size_bytes
1684
% self._physical_blocks_per_logical_kv_block
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py
+13
-1
@@ -75,8 +75,14 @@ class MooncakeStoreCoordinator:
75
retention_interval: int | None = None,
76
dcp_world_size: int = 1,
77
) -> None:
78
+ # Mirrors core's resolve_kv_cache_block_sizes: the hash unit only has
79
+ # to divide groups that participate in prefix caching. Non-shareable
80
+ # scratch groups (e.g. GLM-5.3-Flash's kpool tail) are skipped by
81
+ # _verify_and_split_kv_cache_groups and never probed for hits.
82
assert all(
79
- g.kv_cache_spec.block_size % hash_block_size == 0 for g in kv_cache_groups
83
+ g.kv_cache_spec.block_size % hash_block_size == 0
84
+ for g in kv_cache_groups
85
+ if g.kv_cache_spec.prefix_cacheable
86
), "block_size must be divisible by hash_block_size"
87
assert scheduler_block_size % hash_block_size == 0, (
88
f"scheduler_block_size ({scheduler_block_size}) must be a multiple of "
@@ -116,6 +122,12 @@ class MooncakeStoreCoordinator:
122
"""
123
attention_groups: list[SpecGroup] = []
124
for i, g in enumerate(self.kv_cache_groups):
125
+ # Skip groups that opt out of prefix caching (e.g. GLM-5.3-Flash
126
+ # kpool tail): per-request scratch, never shareable, so they must
127
+ # not participate in hit lookup. Mirrors core's
128
+ # KVCacheCoordinator.verify_and_split_kv_cache_groups.
129
+ if not g.kv_cache_spec.prefix_cacheable:
130
+ continue
131
spec = _unwrap_spec(g.kv_cache_spec)
132
manager_cls = KVCacheSpecRegistry.get_manager_class(spec)
133
assert manager_cls is not None, (
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
+38
-2
@@ -523,6 +523,7 @@ class KVCacheStoreSendingThread(KVTransferThread):
523
enable_group_semantics: bool = False,
524
supports_group_ids: bool = False,
525
record_operation: Callable[..., None] | None = None,
526
+ group_participates: Sequence[bool] | None = None,
527
):
528
super().__init__(
529
store,
@@ -537,6 +538,11 @@ class KVCacheStoreSendingThread(KVTransferThread):
538
self.group_put_steps = group_put_steps
539
self.coord = coord
540
self.kv_role = kv_role
541
+ self.group_participates = (
542
+ list(group_participates)
543
+ if group_participates is not None
544
+ else [True] * len(token_databases)
545
+ )
546
# req_id -> ids of its store jobs that are still queued or running.
547
# Keying by store_job_id, which never repeats for the engine's lifetime,
548
# rather than counting jobs per request id makes the ledger immune to id
@@ -744,6 +750,8 @@ class KVCacheStoreSendingThread(KVTransferThread):
750
saved = self._saved_offset.get(req_meta.req_id, 0)
751
puts: list[tuple[str, list[int], list[int], KeyMetadata]] = []
752
for g_idx, db in enumerate(self.token_databases):
753
+ if not self.group_participates[g_idx]:
754
+ continue
755
group_blocks = req_meta.block_ids[g_idx]
756
# Distribute across ranks by the same rule as normal chunks.
757
put_step = self.group_put_steps[g_idx]
@@ -983,6 +991,8 @@ class KVCacheStoreSendingThread(KVTransferThread):
991
group_indices: list[int] = []
992
store_shard_ids: list[StoreShardId] = []
993
for g_idx, db in enumerate(self.token_databases):
994
+ if not self.group_participates[g_idx]:
995
+ continue
996
# Rotate the stride phase per group to balance load across ranks.
997
put_step = self.group_put_steps[g_idx]
998
put_step_rank = (self.tp_rank + g_idx) % put_step
@@ -1253,6 +1263,7 @@ class KVCacheStoreRecvingThread(KVTransferThread):
1263
disk_offload_buffer_budget_bytes: int | None = None,
1264
record_operation: Callable[..., None] | None = None,
1265
request_queue: queue.Queue[Any] | None = None,
1266
+ group_participates: Sequence[bool] | None = None,
1267
):
1268
super().__init__(
1269
store,
@@ -1264,6 +1275,11 @@ class KVCacheStoreRecvingThread(KVTransferThread):
1275
record_operation=record_operation,
1276
request_queue=request_queue,
1277
)
1278
+ self.group_participates = (
1279
+ list(group_participates)
1280
+ if group_participates is not None
1281
+ else [True] * len(token_databases)
1282
+ )
1283
# _invalid_block_ids can be access by both the Worker and RecvingThread
1284
self._invalid_block_ids_lock = threading.Lock()
1285
self._invalid_block_ids: set[int] = set()
@@ -1311,6 +1327,8 @@ class KVCacheStoreRecvingThread(KVTransferThread):
1327
key_list: list[str] = []
1328
block_id_list: list[int] = []
1329
for g_idx, db in enumerate(self.token_databases):
1330
+ if not self.group_participates[g_idx]:
1331
+ continue
1332
mask = load_mask_per_group[g_idx]
1333
chunks: list[tuple[int, int]] = []
1334
store_shard_ids: list[StoreShardId] = []
@@ -1791,6 +1809,11 @@ class MooncakeStoreWorker:
1809
"""Construct token databases and their Store layouts."""
1810
token_dbs: list[ChunkedTokenDatabase] = []
1811
for group_idx, group in enumerate(self._kv_cache_groups):
1812
+ hash_block_size = (
1813
+ self.hash_block_size
1814
+ if group.kv_cache_spec.prefix_cacheable
1815
+ else group.kv_cache_spec.block_size
1816
+ )
1817
group_tp_rank = self.tp_rank
1818
if layout_cls is None:
1819
group_tp_rank //= self._group_tp_replication_factors[group_idx]
@@ -1805,7 +1828,7 @@ class MooncakeStoreWorker:
1828
store_layout = layout_cls(
1829
group_metadata,
1830
group.kv_cache_spec.block_size,
1808
- self.hash_block_size,
1831
+ hash_block_size,
1832
local_tp_size=self.tp_size,
1833
store_tp_size=self.store_tp_size,
1834
tp_rank=self.tp_rank,
@@ -1815,7 +1838,7 @@ class MooncakeStoreWorker:
1838
ChunkedTokenDatabase(
1839
group_metadata,
1840
group.kv_cache_spec.block_size,
1818
- hash_block_size=self.hash_block_size,
1841
+ hash_block_size=hash_block_size,
1842
store_layout=store_layout,
1843
)
1844
)
@@ -1937,6 +1960,10 @@ class MooncakeStoreWorker:
1960
enable_group_semantics=self.enable_group_semantics,
1961
supports_group_ids=self._supports_group_ids,
1962
record_operation=self._record_kv_connector_operation,
1963
+ group_participates=[
1964
+ group.kv_cache_spec.prefix_cacheable
1965
+ for group in self._kv_cache_groups
1966
+ ],
1967
)
1968
self.kv_send_thread.start()
1969
@@ -1954,6 +1981,10 @@ class MooncakeStoreWorker:
1981
disk_offload_buffer_budget_bytes=self.disk_offload_buffer_budget_bytes,
1982
record_operation=self._record_kv_connector_operation,
1983
request_queue=self.recv_request_queue,
1984
+ group_participates=[
1985
+ group.kv_cache_spec.prefix_cacheable
1986
+ for group in self._kv_cache_groups
1987
+ ],
1988
)
1989
recv_thread.name = f"KVCacheStoreRecvingThread-{i}"
1990
recv_thread.start()
@@ -2126,6 +2157,8 @@ class MooncakeStoreWorker:
2157
fine_grained = self.coord.enable_partial_hash_hits
2158
lookup_masks = None if fine_grained else self.coord.lookup_mask(token_len)
2159
for g_idx, db in enumerate(self.token_dbs):
2160
+ if not self._kv_cache_groups[g_idx].kv_cache_spec.prefix_cacheable:
2161
+ continue
2162
spec_block_size = db.block_size
2163
key_prefixes = self._lookup_key_prefixes[g_idx]
2164
if fine_grained:
@@ -2236,6 +2269,9 @@ class MooncakeStoreWorker:
2269
boundaries = []
2270
hit_boundary_hash_idx = hit_length // self.hash_block_size - 1
2271
for group_id, db in enumerate(self.token_dbs):
2272
+ if not self._kv_cache_groups[group_id].kv_cache_spec.prefix_cacheable:
2273
+ # Scratch groups are never stored, so they have no tail key.
2274
+ continue
2275
chunk_id = cdiv(hit_length, db.block_size) - 1
2276
boundary_tokens = hit_length
2277
contains_hit_boundary = cached_block_pool.contains(
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+112
-6
@@ -74,6 +74,7 @@ from vllm.utils.torch_utils import async_tensor_h2d
74
from vllm.v1.kv_cache_interface import (
75
CircularBufferSpec,
76
FullAttentionSpec,
77
+ KpoolTailSpec,
78
KVCacheLayout,
79
KVCacheSpec,
80
MambaSpec,
@@ -99,6 +100,46 @@ def _share_storage_and_block_stride(caches: list[torch.Tensor]) -> bool:
100
return len(block_strides) == len(storage_ptrs) == 1
101
102
103
+def _tensor_byte_span_end(cache: torch.Tensor) -> int:
104
+ """Return the exclusive end address touched by a nonnegative-stride view."""
105
+ if cache.numel() == 0:
106
+ return cache.data_ptr()
107
+ if any(stride < 0 for stride in cache.stride()):
108
+ raise ValueError("NIXL cache views must have nonnegative strides")
109
+ max_element_offset = sum(
110
+ (size - 1) * stride for size, stride in zip(cache.shape, cache.stride())
111
+ )
112
+ return cache.data_ptr() + (max_element_offset + 1) * cache.element_size()
113
+
114
+
115
+def _uses_dense_virtual_transfer_pages(
116
+ layer_spec: KVCacheSpec,
117
+ cache: torch.Tensor,
118
+ physical_page_size: int,
119
+ num_blocks: int,
120
+) -> bool:
121
+ """Return whether a compressed kernel view can be split into NIXL pages."""
122
+ if not (
123
+ isinstance(layer_spec, MLAAttentionSpec)
124
+ and layer_spec.tokens_per_state > 1
125
+ and cache.ndim == 4
126
+ and cache.shape[1] == 1
127
+ and cache.is_contiguous()
128
+ and physical_page_size > 0
129
+ and layer_spec.state_content_size_bytes > 0
130
+ ):
131
+ return False
132
+
133
+ block_stride = cache.stride(0) * cache.element_size()
134
+ return (
135
+ block_stride > physical_page_size
136
+ and block_stride % physical_page_size == 0
137
+ and physical_page_size % layer_spec.state_content_size_bytes == 0
138
+ and cache.shape[0] * (block_stride // physical_page_size) == num_blocks
139
+ and cache.nbytes == num_blocks * physical_page_size
140
+ )
141
+
142
+
143
class NixlBaseConnectorWorker:
144
"""Base implementation of Worker side methods shared by pull and push."""
145
@@ -1150,6 +1191,31 @@ class NixlBaseConnectorWorker:
1191
# aliases even though every view shares one block-major allocation.
1192
packed_storage = packed_storage and not self._is_csa_linear
1193
1194
+ layer_specs: dict[str, KVCacheSpec] = {}
1195
+ compressed_region_owners: dict[int, torch.Tensor] = {}
1196
+ for layer_name, cache in xfer_buffers.items():
1197
+ layer_spec = self._layer_specs.get(layer_name)
1198
+ if isinstance(layer_spec, UniformTypeKVCacheSpecs):
1199
+ layer_spec = layer_spec.kv_cache_specs[layer_name]
1200
+ if layer_spec is None:
1201
+ continue
1202
+ layer_specs[layer_name] = layer_spec
1203
+ physical_page_size = (
1204
+ layer_spec.page_size_bytes
1205
+ if isinstance(layer_spec, MambaSpec)
1206
+ else layer_spec.page_size_bytes
1207
+ // self._physical_blocks_per_logical_kv_block
1208
+ )
1209
+ num_blocks = (
1210
+ self._logical_num_blocks
1211
+ if isinstance(layer_spec, MambaSpec)
1212
+ else self.num_blocks
1213
+ )
1214
+ if _uses_dense_virtual_transfer_pages(
1215
+ layer_spec, cache, physical_page_size, num_blocks
1216
+ ):
1217
+ compressed_region_owners.setdefault(cache.data_ptr(), cache)
1218
+
1219
# K and V are packed into the content dim, so each attention layer is a
1220
# single NIXL region whose block transfers as one unit. Mamba layers instead
1221
# register separate conv/ssm sub-regions (see `_build_mamba_local`).
@@ -1158,7 +1224,7 @@ class NixlBaseConnectorWorker:
1224
# kernel requires a specific block size. This leads to SSM and FA layers
1225
# having different num_blocks.
1226
# `_physical_blocks_per_logical_kv_block` ratio is used to adjust for this.
1161
- layer_spec = self._layer_specs.get(layer_name)
1227
+ layer_spec = layer_specs.get(layer_name)
1228
if layer_spec is None:
1229
logger.debug(
1230
"Skipping layer %s as no KVCache spec is present. "
@@ -1166,9 +1232,6 @@ class NixlBaseConnectorWorker:
1232
layer_name,
1233
)
1234
continue
1169
- if isinstance(layer_spec, UniformTypeKVCacheSpecs):
1170
- # DSA Indexer case: UniformTypeKVCacheSpecs merges kv_cache_specs
1171
- layer_spec = layer_spec.kv_cache_specs[layer_name]
1235
# `layer_spec.page_size_bytes` only accounts for logical page_size, that is
1236
# the page_size assuming constant `self._logical_num_blocks`.
1237
physical_page_size = (
@@ -1187,6 +1250,33 @@ class NixlBaseConnectorWorker:
1250
)
1251
storage = cache.untyped_storage()
1252
storage_addr = storage.data_ptr()
1253
+
1254
+ if isinstance(layer_spec, KpoolTailSpec):
1255
+ compressed_owner = compressed_region_owners.get(cache.data_ptr())
1256
+ if compressed_owner is not None:
1257
+ owner_storage = compressed_owner.untyped_storage()
1258
+ owner_end = compressed_owner.data_ptr() + compressed_owner.nbytes
1259
+ tail_is_covered = (
1260
+ compressed_owner.is_contiguous()
1261
+ and owner_storage.data_ptr() == storage_addr
1262
+ and _tensor_byte_span_end(cache) <= owner_end
1263
+ )
1264
+ if not tail_is_covered:
1265
+ raise AssertionError(
1266
+ "Kpool tail cache is not fully covered by its compressed "
1267
+ f"indexer region: layer={layer_name}, "
1268
+ f"tail_shape={tuple(cache.shape)}, "
1269
+ f"tail_stride={tuple(cache.stride())}, "
1270
+ f"owner_shape={tuple(compressed_owner.shape)}, "
1271
+ f"owner_stride={tuple(compressed_owner.stride())}"
1272
+ )
1273
+ logger.debug(
1274
+ "Skipping layer %s because its compressed indexer region "
1275
+ "covers the same storage",
1276
+ layer_name,
1277
+ )
1278
+ continue
1279
+
1280
# Memory registration follows allocations, while transfer regions follow
1281
# logical layers (or contiguous head segments). This keeps strided
1282
# cross-layer views inside their registered allocation.
@@ -1228,7 +1318,15 @@ class NixlBaseConnectorWorker:
1318
and cache.stride(2) == cache.shape[3]
1319
and cache.stride(1) == cache.shape[2] * cache.shape[3]
1320
)
1231
- if storage_is_block_major and (
1321
+ virtual_transfer_pages = _uses_dense_virtual_transfer_pages(
1322
+ layer_spec, cache, physical_page_size, num_blocks
1323
+ )
1324
+ if virtual_transfer_pages:
1325
+ # A compressed kernel row can contain multiple NIXL transfer pages.
1326
+ region_specs = [
1327
+ (cache.data_ptr(), physical_page_size, physical_page_size)
1328
+ ]
1329
+ elif storage_is_block_major and (
1330
(packed_storage and is_mla_region)
1331
or (not hnc_contiguous and not self._is_csa_linear)
1332
):
@@ -1243,7 +1341,15 @@ class NixlBaseConnectorWorker:
1341
]
1342
else:
1343
segment_bytes = num_blocks * block_stride
1246
- assert cache.nbytes % segment_bytes == 0
1344
+ if cache.nbytes % segment_bytes != 0:
1345
+ raise AssertionError(
1346
+ "KV cache view cannot be partitioned into NIXL regions: "
1347
+ f"layer={layer_name}, cache_nbytes={cache.nbytes}, "
1348
+ f"num_blocks={num_blocks}, block_stride={block_stride}, "
1349
+ f"physical_page_size={physical_page_size}, "
1350
+ f"cache_shape={tuple(cache.shape)}, "
1351
+ f"cache_stride={tuple(cache.stride())}"
1352
+ )
1353
num_segments = cache.nbytes // segment_bytes
1354
region_block_len = (
1355
block_stride if num_segments > 1 else physical_page_size
vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py
+2
-1
@@ -45,8 +45,9 @@ PUSH_REG_NOTIF_PREFIX = b"PUSH_REG:"
45
# 7: Include NIXL transfer mode (push vs pull) in the compatibility hash
46
# 8: Add dcp_size and pcp_size to NixlAgentMetadata
47
# 9: Add block_strides
48
+# 10: Add dense virtual transfer pages for compressed MLA caches
49
#
49
-NIXL_CONNECTOR_VERSION: int = 9
50
+NIXL_CONNECTOR_VERSION: int = 10
51
52
53
@dataclass
vllm/model_executor/layers/attention/mla_attention.py
+1
-1
@@ -1534,7 +1534,7 @@ class MLACommonBackend(AttentionBackend):
1534
1535
@classmethod
1536
def get_supported_head_sizes(cls) -> list[int]:
1537
- return [320, 576]
1537
+ return [320, 512, 576]
1538
1539
@classmethod
1540
def is_mla(cls) -> bool:
vllm/model_executor/layers/mamba/ops/scatter_states.py
new
+72
@@ -0,0 +1,72 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+
6
+from vllm.platforms import current_platform
7
+from vllm.triton_utils import tl, triton
8
+
9
+
10
+@triton.jit
11
+def _scatter_states_kernel(
12
+ state_ptr,
13
+ src_ptr,
14
+ indices_ptr,
15
+ stride_state_batch,
16
+ stride_src_batch,
17
+ stride_indices,
18
+ row_size: tl.constexpr,
19
+ BLOCK_SIZE: tl.constexpr,
20
+ launch_pdl: tl.constexpr,
21
+):
22
+ block_idx = tl.program_id(0)
23
+ batch_idx = tl.program_id(1)
24
+ offsets = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
25
+ mask = offsets < row_size
26
+
27
+ if launch_pdl:
28
+ tl.extra.cuda.gdc_wait()
29
+ tl.extra.cuda.gdc_launch_dependents()
30
+
31
+ state_idx = tl.load(indices_ptr + batch_idx * stride_indices).to(tl.int64)
32
+ values = tl.load(src_ptr + batch_idx * stride_src_batch + offsets, mask=mask)
33
+ tl.store(state_ptr + state_idx * stride_state_batch + offsets, values, mask=mask)
34
+
35
+
36
+def scatter_states(
37
+ state: torch.Tensor,
38
+ src: torch.Tensor,
39
+ indices: torch.Tensor,
40
+) -> None:
41
+ """Scatter ``src`` rows into ``state`` at ``indices`` (in place).
42
+
43
+ Equivalent to ``state[indices] = src`` but non-atomic and bandwidth-bound,
44
+ since mamba cache slots are unique per sequence. ``gather_initial_states``
45
+ is the read-side counterpart.
46
+ """
47
+ assert state.ndim >= 2
48
+ assert state.is_cuda
49
+ assert src.ndim == state.ndim
50
+ assert indices.ndim == 1
51
+ assert indices.device == state.device
52
+ assert src.shape[1:] == state.shape[1:]
53
+ assert src.shape[0] == indices.shape[0]
54
+ assert indices.dtype in (torch.int32, torch.int64)
55
+
56
+ row_size = state[0].numel()
57
+ assert state[0].is_contiguous()
58
+ assert src[0].is_contiguous()
59
+ block_size = min(triton.next_power_of_2(row_size), 1024)
60
+ grid = (triton.cdiv(row_size, block_size), indices.numel())
61
+ _scatter_states_kernel[grid](
62
+ state,
63
+ src,
64
+ indices,
65
+ state.stride(0),
66
+ src.stride(0),
67
+ indices.stride(0),
68
+ row_size=row_size,
69
+ BLOCK_SIZE=block_size,
70
+ num_warps=8,
71
+ launch_pdl=current_platform.is_arch_support_pdl(),
72
+ )
vllm/model_executor/layers/mhc.py
+97
-11
@@ -1,5 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import inspect
4
+
5
import torch
6
7
# this import will also register the custom ops
@@ -29,6 +31,67 @@ HAS_TILELANG_MHC = _has_tilelang_mhc()
31
HAS_AITER_MHC = is_aiter_found_and_supported()
32
33
34
+def _has_aiter_mhc_fused() -> bool:
35
+ if not HAS_AITER_MHC:
36
+ return False
37
+ try:
38
+ from aiter.ops.mhc import mhc_fused_post_pre
39
+ except Exception:
40
+ return False
41
+ return callable(mhc_fused_post_pre)
42
+
43
+
44
+def _aiter_mhc_op_accepts_norm(op_name: str) -> bool:
45
+ if not HAS_AITER_MHC:
46
+ return False
47
+ try:
48
+ from aiter.ops import mhc as aiter_mhc
49
+
50
+ for candidate_name in (op_name, f"{op_name}_fake"):
51
+ op = getattr(aiter_mhc, candidate_name, None)
52
+ if op is None:
53
+ continue
54
+ parameters = inspect.signature(op).parameters
55
+ if "norm_weight" in parameters and "norm_eps" in parameters:
56
+ return True
57
+ except (AttributeError, ImportError, TypeError, ValueError):
58
+ pass
59
+ return False
60
+
61
+
62
+HAS_AITER_MHC_FUSED = _has_aiter_mhc_fused()
63
+HAS_AITER_MHC_PRE_NORM = _aiter_mhc_op_accepts_norm("mhc_pre")
64
+HAS_AITER_MHC_FUSED_NORM = _aiter_mhc_op_accepts_norm("mhc_fused_post_pre")
65
+
66
+
67
+def _aiter_mhc_supported(
68
+ residual: torch.Tensor,
69
+ norm_weight: torch.Tensor | None,
70
+ *,
71
+ supports_norm: bool,
72
+) -> bool:
73
+ hidden_size = residual.shape[-1]
74
+ hc_mult = residual.shape[-2]
75
+ return (
76
+ HAS_AITER_MHC
77
+ and hidden_size % 256 == 0
78
+ and hc_mult == 4
79
+ and (norm_weight is None or supports_norm)
80
+ )
81
+
82
+
83
+def _apply_mhc_norm(
84
+ layer_input: torch.Tensor,
85
+ norm_weight: torch.Tensor | None,
86
+ norm_eps: float,
87
+) -> torch.Tensor:
88
+ if norm_weight is None:
89
+ return layer_input
90
+ from vllm import ir
91
+
92
+ return ir.ops.rms_norm(layer_input, norm_weight, norm_eps)
93
+
94
+
95
# --8<-- [start:mhc_pre]
96
@CustomOp.register("mhc_pre")
97
class MHCPreOp(CustomOp):
@@ -89,9 +152,11 @@ class MHCPreOp(CustomOp):
152
norm_weight: torch.Tensor | None = None,
153
norm_eps: float = 0.0,
154
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
92
- hidden_size = residual.shape[-1]
93
- hc_mult = residual.shape[-2]
94
- if HAS_AITER_MHC and hidden_size % 256 == 0 and hc_mult == 4:
155
+ if _aiter_mhc_supported(
156
+ residual,
157
+ norm_weight,
158
+ supports_norm=HAS_AITER_MHC_PRE_NORM,
159
+ ):
160
return torch.ops.vllm.mhc_pre_aiter(
161
residual,
162
fn,
@@ -122,7 +187,7 @@ class MHCPreOp(CustomOp):
187
norm_eps,
188
)
189
else:
125
- return self.forward_native(
190
+ post_mix, comb_mix, layer_input = self.forward_native(
191
residual,
192
fn,
193
hc_scale,
@@ -136,6 +201,11 @@ class MHCPreOp(CustomOp):
201
norm_weight,
202
norm_eps,
203
)
204
+ return (
205
+ post_mix,
206
+ comb_mix,
207
+ _apply_mhc_norm(layer_input, norm_weight, norm_eps),
208
+ )
209
210
def forward_native(
211
self,
@@ -225,9 +295,7 @@ class MHCPostOp(CustomOp):
295
post_layer_mix: torch.Tensor,
296
comb_res_mix: torch.Tensor,
297
) -> torch.Tensor:
228
- hidden_size = residual.shape[-1]
229
- hc_mult = residual.shape[-2]
230
- if HAS_AITER_MHC and hidden_size % 256 == 0 and hc_mult == 4:
298
+ if _aiter_mhc_supported(residual, None, supports_norm=True):
299
return torch.ops.vllm.mhc_post_aiter(
300
x,
301
residual,
@@ -449,9 +517,11 @@ class MHCFusedPostPreOp(CustomOp):
517
norm_weight: torch.Tensor | None = None,
518
norm_eps: float = 0.0,
519
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
452
- hidden_size = residual.shape[-1]
453
- hc_mult = residual.shape[-2]
454
- if HAS_AITER_MHC and hidden_size % 256 == 0 and hc_mult == 4:
520
+ if HAS_AITER_MHC_FUSED and _aiter_mhc_supported(
521
+ residual,
522
+ norm_weight,
523
+ supports_norm=HAS_AITER_MHC_FUSED_NORM,
524
+ ):
525
return torch.ops.vllm.mhc_fused_post_pre_aiter(
526
x,
527
residual,
@@ -489,7 +559,7 @@ class MHCFusedPostPreOp(CustomOp):
559
norm_weight,
560
norm_eps,
561
)
492
- return self.forward_native(
562
+ residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur = self.forward_native(
563
x,
564
residual,
565
post_layer_mix,
@@ -507,6 +577,12 @@ class MHCFusedPostPreOp(CustomOp):
577
norm_weight,
578
norm_eps,
579
)
580
+ return (
581
+ residual_cur,
582
+ post_mix_cur,
583
+ comb_mix_cur,
584
+ _apply_mhc_norm(layer_input_cur, norm_weight, norm_eps),
585
+ )
586
587
def forward_native(
588
self,
@@ -577,3 +653,13 @@ class MHCFusedPostPreOp(CustomOp):
653
hc_post_mult_value,
654
sinkhorn_repeat,
655
)
656
+
657
+
658
+def hc_expand(x: torch.Tensor, n: int) -> torch.Tensor:
659
+ """[s, hidden_size] -> [s, n * hidden_size] by replication."""
660
+ return x.unsqueeze(1).expand(-1, n, -1).contiguous()
661
+
662
+
663
+def hc_contract(x: torch.Tensor, n: int) -> torch.Tensor:
664
+ """[s, n * hidden_size] -> [s, hidden_size] by averaging."""
665
+ return x.mean(dim=1)
vllm/model_executor/layers/mla.py
+22
-3
@@ -8,6 +8,7 @@ from vllm.config import CacheConfig
8
from vllm.model_executor.custom_op import PluggableLayer
9
from vllm.model_executor.layers.attention import MLAAttention
10
from vllm.model_executor.layers.quantization import QuantizationConfig
11
+from vllm.models.common.ops import fused_q_kv_rmsnorm
12
from vllm.platforms import current_platform
13
14
@@ -17,7 +18,7 @@ class MLAModules:
18
19
kv_a_layernorm: torch.nn.Module
20
kv_b_proj: torch.nn.Module
20
- rotary_emb: torch.nn.Module
21
+ rotary_emb: torch.nn.Module | None
22
o_proj: torch.nn.Module
23
fused_qkv_a_proj: torch.nn.Module | None
24
kv_a_proj_with_mqa: torch.nn.Module | None
@@ -69,6 +70,7 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):
70
skip_topk: bool = False,
71
non_causal_multi_token_decode: bool = False,
72
allow_short_prefill_indexer_scoring_skip: bool = False,
73
+ fuse_qkv_rmsnorm: bool = False,
74
) -> None:
75
super().__init__()
76
self.hidden_size = hidden_size
@@ -98,6 +100,9 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):
100
# the topk_tokens buffer written by a previous layer in the same pass.
101
# Refer: https://arxiv.org/abs/2603.12201 for more details.
102
self.skip_topk = skip_topk
103
+ # When True, fuse the q_a and kv_a RMSNorms into a single kernel launch
104
+ # (MLA layers with q-LoRA). Opt-in; default False preserves other models.
105
+ self.fuse_qkv_rmsnorm = fuse_qkv_rmsnorm
106
# qrep is active when the query projection is a DCP-group-sharded layer
107
# that materializes the full group head set locally.
108
q_proj_layer = self.q_b_proj if self.q_lora_rank is not None else self.q_proj
@@ -172,8 +177,9 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):
177
[self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
178
dim=-1,
179
)
175
- q_c = self.q_a_layernorm(q_c)
180
q_proj_layer = self.q_b_proj
181
+ if not self.fuse_qkv_rmsnorm:
182
+ q_c = self.q_a_layernorm(q_c)
183
q_proj_input = q_c
184
else:
185
assert self.kv_a_proj_with_mqa is not None, (
@@ -187,7 +193,20 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):
193
q_proj_input = hidden_states
194
195
kv_c, k_pe = kv_lora.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
190
- kv_c_normed = self.kv_a_layernorm(kv_c)
196
+ if self.fuse_qkv_rmsnorm and q_c is not None:
197
+ # Fuse q_a + kv_a RMSNorm into one kernel launch (q_c is still the
198
+ # pre-norm projection output here).
199
+ assert self.q_a_layernorm is not None
200
+ q_proj_input, kv_c_normed = fused_q_kv_rmsnorm(
201
+ q_c,
202
+ kv_c,
203
+ self.q_a_layernorm.weight.data,
204
+ self.kv_a_layernorm.weight.data,
205
+ self.q_a_layernorm.variance_epsilon,
206
+ )
207
+ q_c = q_proj_input
208
+ else:
209
+ kv_c_normed = self.kv_a_layernorm(kv_c)
210
# Add head dim of 1 to k_pe
211
k_pe = k_pe.unsqueeze(1)
212
vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
+57
-76
@@ -361,30 +361,27 @@ def _shuffle_deepseek_fp8_moe_weights(
361
Returns 4D weight tensors in BlockMajorK layout
362
(E, K/block_k, Mn, block_k)
363
"""
364
- from flashinfer import shuffle_matrix_a
365
- from flashinfer.fused_moe import convert_to_block_layout
364
+ from flashinfer.utils import get_shuffle_matrix_a_row_indices
365
366
epilogue_tile_m = 64
367
block_k = 128
369
- num_experts = w13.shape[0]
368
371
- M13, K13 = w13.shape[1], w13.shape[2]
372
- M2, K2 = w2.shape[1], w2.shape[2]
373
- w13_out = torch.empty(
374
- num_experts, K13 // block_k, M13, block_k, dtype=torch.uint8, device=w13.device
375
- )
376
- w2_out = torch.empty(
377
- num_experts, K2 // block_k, M2, block_k, dtype=torch.uint8, device=w2.device
378
- )
379
-
380
- for i in range(num_experts):
381
- t13 = shuffle_matrix_a(w13[i].view(torch.uint8), epilogue_tile_m)
382
- w13_out[i] = convert_to_block_layout(t13, block_k)
383
-
384
- t2 = shuffle_matrix_a(w2[i].view(torch.uint8), epilogue_tile_m)
385
- w2_out[i] = convert_to_block_layout(t2, block_k)
369
+ def shuffle_to_block_major_k(w: torch.Tensor) -> torch.Tensor:
370
+ # shuffle_matrix_a's row permutation depends only on (M,
371
+ # epilogue_tile_m), so it is computed once and applied to every expert
372
+ # in a single gather instead of once per expert. Gathering through the
373
+ # BlockMajorK-permuted view also folds convert_to_block_layout into
374
+ # that same kernel. Per-expert loops here cost minutes for a MoE this
375
+ # wide (~24k tiny launches plus a host round-trip each).
376
+ num_experts, m, k = w.shape
377
+ rows = get_shuffle_matrix_a_row_indices(
378
+ w[0].view(torch.uint8), epilogue_tile_m
379
+ ).to(w.device)
380
+ blocked = w.view(torch.uint8).view(num_experts, m, k // block_k, block_k)
381
+ out = blocked.permute(0, 2, 1, 3)[:, :, rows, :].contiguous()
382
+ return out.view(torch.float8_e4m3fn)
383
387
- return w13_out.view(torch.float8_e4m3fn), w2_out.view(torch.float8_e4m3fn)
384
+ return shuffle_to_block_major_k(w13), shuffle_to_block_major_k(w2)
385
386
387
def _shuffle_mxfp8_moe_weights(
@@ -396,70 +393,54 @@ def _shuffle_mxfp8_moe_weights(
393
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
394
"""Preprocess MXFP8 weights and scales for the FlashInfer TRT-LLM kernel.
395
399
- Following flashinfer/tests/moe/test_trtllm_gen_fused_moe.py:
400
- 1. reorder_rows_for_gated_act_gemm (interleave gate/up rows)
401
- 2. shuffle_matrix_a (weight data layout shuffle)
402
- 3. shuffle_matrix_sf_a (scale factor layout shuffle)
396
+ All three transforms (gate/up row reorder, ``shuffle_matrix_a`` weight
397
+ shuffle, ``shuffle_matrix_sf_a`` scale shuffle) are fixed row/index
398
+ permutations that depend only on the per-expert matrix shape, so the
399
+ permutation is computed once and applied to every expert in a single
400
+ gather instead of once per expert. ``block_scale_interleave`` accepts a
401
+ batched ``(E, M, K)`` scale tensor directly. Output is bit-identical to
402
+ the per-expert loop but ~20x faster (a 288-expert MoE otherwise costs
403
+ seconds per layer at load).
404
"""
404
- from flashinfer import (
405
- reorder_rows_for_gated_act_gemm,
406
- shuffle_matrix_a,
407
- shuffle_matrix_sf_a,
405
+ from flashinfer.fused_moe.core import (
406
+ get_reorder_rows_for_gated_act_gemm_row_indices,
407
)
408
+ from flashinfer.quantization.fp4_quantization import block_scale_interleave
409
+ from flashinfer.utils import get_shuffle_matrix_a_row_indices
410
411
epilogue_tile_m = 128
411
- num_experts = w13.shape[0]
412
- intermediate_size = w13.shape[1] // 2
413
- hidden_size = w13.shape[2]
412
415
- w13_interleaved: list[torch.Tensor] = []
416
- w13_scale_interleaved: list[torch.Tensor] = []
417
- for i in range(num_experts):
418
- if is_gated:
419
- w13_interleaved.append(
420
- reorder_rows_for_gated_act_gemm(
421
- w13[i].reshape(2 * intermediate_size, -1)
422
- )
423
- )
424
- w13_scale_interleaved.append(
425
- reorder_rows_for_gated_act_gemm(
426
- w13_scale[i].reshape(2 * intermediate_size, -1)
427
- )
428
- )
429
- else:
430
- w13_interleaved.append(w13[i])
431
- w13_scale_interleaved.append(w13_scale[i])
432
-
433
- w13_shuffled: list[torch.Tensor] = []
434
- w2_shuffled: list[torch.Tensor] = []
435
- w13_scale_shuffled: list[torch.Tensor] = []
436
- w2_scale_shuffled: list[torch.Tensor] = []
437
- for i in range(num_experts):
438
- w13_shuffled.append(
439
- shuffle_matrix_a(w13_interleaved[i].view(torch.uint8), epilogue_tile_m)
440
- )
441
- w2_shuffled.append(shuffle_matrix_a(w2[i].view(torch.uint8), epilogue_tile_m))
442
- w13_scale_shuffled.append(
443
- shuffle_matrix_sf_a(
444
- w13_scale_interleaved[i]
445
- .view(torch.uint8)
446
- .reshape(2 * intermediate_size, -1),
447
- epilogue_tile_m,
448
- )
449
- )
450
- w2_scale_shuffled.append(
451
- shuffle_matrix_sf_a(
452
- w2_scale[i].view(torch.uint8).reshape(hidden_size, -1),
453
- epilogue_tile_m,
454
- )
455
- )
413
+ w13_u = w13.view(torch.uint8)
414
+ w2_u = w2.view(torch.uint8)
415
457
- w13_out = torch.stack(w13_shuffled).view(torch.float8_e4m3fn)
458
- w2_out = torch.stack(w2_shuffled).view(torch.float8_e4m3fn)
459
- w13_scale_out = torch.stack(w13_scale_shuffled).reshape(w13_scale.shape)
460
- w2_scale_out = torch.stack(w2_scale_shuffled).reshape(w2_scale.shape)
416
+ # 1. Interleave gate/up rows (gated activation GEMM layout).
417
+ if is_gated:
418
+ gate_idx = get_reorder_rows_for_gated_act_gemm_row_indices(
419
+ w13_u[0].reshape(w13_u.shape[1], -1)
420
+ )
421
+ w13_u = w13_u[:, gate_idx]
422
+ w13_scale = w13_scale[:, gate_idx]
423
+
424
+ def shuffle_weights(t: torch.Tensor) -> torch.Tensor:
425
+ # Row permutation depends only on (M, epilogue_tile_m).
426
+ idx = get_shuffle_matrix_a_row_indices(t[0], epilogue_tile_m).to(t.device)
427
+ return t[:, idx].view(torch.float8_e4m3fn)
428
+
429
+ def shuffle_scales(s: torch.Tensor) -> torch.Tensor:
430
+ # shuffle_matrix_sf_a == row-gather (same indices as the weight shuffle)
431
+ # followed by the 128x4 block-scale interleave, which is batch-capable.
432
+ idx = get_shuffle_matrix_a_row_indices(
433
+ s[0].view(torch.uint8).reshape(s.shape[1], -1), epilogue_tile_m
434
+ ).to(s.device)
435
+ interleaved = block_scale_interleave(s[:, idx])
436
+ return interleaved.reshape(s.shape)
437
462
- return w13_out, w2_out, w13_scale_out, w2_scale_out
438
+ return (
439
+ shuffle_weights(w13_u),
440
+ shuffle_weights(w2_u),
441
+ shuffle_scales(w13_scale),
442
+ shuffle_scales(w2_scale),
443
+ )
444
445
446
def prepare_fp8_moe_layer_for_fi(
vllm/model_executor/layers/sparse_attn_indexer_kpool.py
new
+1053
@@ -0,0 +1,1053 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Custom Sparse Attention Indexer layers."""
4
+
5
+from typing import TYPE_CHECKING
6
+
7
+import torch
8
+
9
+import vllm.envs as envs
10
+from vllm._aiter_ops import rocm_aiter_ops
11
+from vllm.compilation.breakable_cudagraph import eager_break_during_capture
12
+from vllm.config import get_current_vllm_config_or_none
13
+from vllm.forward_context import get_forward_context
14
+from vllm.logger import init_logger
15
+from vllm.model_executor.custom_op import CustomOp
16
+from vllm.platforms import current_platform
17
+
18
+if TYPE_CHECKING:
19
+ from vllm.models.glm5next.nvidia.ops import kpool_compress as kpool_ops
20
+elif current_platform.is_rocm():
21
+ from vllm.models.glm5next.amd.ops import kpool_compress as kpool_ops
22
+else:
23
+ from vllm.models.glm5next.nvidia.ops import kpool_compress as kpool_ops
24
+
25
+from vllm.utils.deep_gemm import has_deep_gemm
26
+from vllm.utils.torch_utils import (
27
+ LayerNameType,
28
+ _encode_layer_name,
29
+ _resolve_layer_name,
30
+)
31
+from vllm.v1.attention.backends.mla.indexer import (
32
+ DeepseekV32IndexerMetadata,
33
+)
34
+from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton
35
+from vllm.v1.worker.workspace import current_workspace_manager
36
+
37
+if current_platform.is_cuda_alike():
38
+ from vllm import _custom_ops as ops
39
+elif current_platform.is_xpu():
40
+ from vllm._xpu_ops import xpu_ops
41
+
42
+logger = init_logger(__name__)
43
+
44
+RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024
45
+
46
+# MXFP4 layout: 2 values packed per byte, ue8m0 (1-byte) scale per block of 32.
47
+MXFP4_BLOCK_SIZE = 32
48
+
49
+# kpool write helper: form pools from the current token batch and compress them
50
+# into the index K cache via the fused Triton kernel.
51
+
52
+
53
+def _kpool_compress_insert(
54
+ k: torch.Tensor,
55
+ gate_score: torch.Tensor,
56
+ ape: torch.Tensor,
57
+ kv_cache: torch.Tensor,
58
+ slot_mapping: torch.Tensor,
59
+ kpool: int,
60
+ head_dim: int,
61
+ round_scale: bool,
62
+) -> None:
63
+ """Pool ``kpool`` consecutive tokens into one fp8 K and write at pool slots.
64
+
65
+ ``slot_mapping`` is pool-granular (compress_ratio == kpool on the spec):
66
+ only the *last* token of each complete pool carries a valid (>=0) slot;
67
+ intra-pool tokens are -1. Every position is treated as a pool-completion
68
+ candidate and non-completions are masked off inside the kernel. Compacting
69
+ the valid rows first costs two device syncs on the eager prefill path and
70
+ buys nothing numerically. Assumes pool-aligned chunk starts.
71
+ """
72
+ n = slot_mapping.shape[0]
73
+ # No pool can complete in a batch smaller than one pool; also keeps the
74
+ # clamped gather indices below in bounds.
75
+ if n < kpool:
76
+ return
77
+ pos = torch.arange(n, device=k.device)
78
+ valid = slot_mapping >= 0
79
+ # Drop pools whose start falls before the batch (leading padding); their
80
+ # gate/k data is undefined anyway.
81
+ write_mask = valid & (pos >= kpool - 1)
82
+ offs = torch.arange(kpool, device=k.device)
83
+ idx = (pos - (kpool - 1)).clamp_min(0)[:, None] + offs[None, :]
84
+ kpool_ops.kpool_compress_and_write_cache(
85
+ kv_cache,
86
+ k[idx], # [n, kpool, head_dim]
87
+ gate_score[idx],
88
+ ape,
89
+ slot_mapping.to(torch.int64),
90
+ pool_size=kpool,
91
+ head_dim=head_dim,
92
+ write_mask=write_mask,
93
+ round_scale=round_scale,
94
+ write_cache=True,
95
+ return_compressed=False,
96
+ )
97
+
98
+
99
+def _build_decode_scatter_indices(
100
+ decode_lens: torch.Tensor,
101
+ num_requests: int,
102
+ n: int,
103
+) -> tuple[torch.Tensor, torch.Tensor]:
104
+ """Per-token (request id, intra-request index) for a non-uniform decode
105
+ batch, with ``n == decode_lens.sum()`` as a host int (avoids a
106
+ device sync and keeps both repeat_interleaves sync-free).
107
+
108
+ Shared by every ``_scatter_decode_tokens_by_request`` call in a step:
109
+ building it per call would repeat the same repeat_interleave/cumsum chain
110
+ up to 5x per layer on the eager decode break.
111
+ """
112
+ device = decode_lens.device
113
+ dl = decode_lens.to(torch.int64)
114
+ req_id = torch.repeat_interleave(
115
+ torch.arange(num_requests, device=device, dtype=torch.int64),
116
+ dl,
117
+ output_size=n,
118
+ )
119
+ req_starts = torch.cumsum(
120
+ torch.cat([torch.zeros(1, device=device, dtype=torch.int64), dl[:-1]]),
121
+ dim=0,
122
+ )
123
+ # Broadcast the per-request start offsets to per-token (length n ==
124
+ # dl.sum()) so each token's intra-request index subtracts its own
125
+ # request's start.
126
+ starts = torch.repeat_interleave(req_starts, dl, output_size=n)
127
+ intra = torch.arange(n, device=device, dtype=torch.int64) - starts
128
+ return req_id, intra
129
+
130
+
131
+def _scatter_decode_tokens_by_request(
132
+ tokens: torch.Tensor,
133
+ pad_value,
134
+ num_requests: int,
135
+ lmax: int,
136
+ scatter_indices: tuple[torch.Tensor, torch.Tensor],
137
+) -> torch.Tensor:
138
+ """Group ``[N, ...]`` decode tokens into a padded ``[num_requests, lmax, ...]``
139
+ layout: request ``r``'s tokens at row ``r`` in order; short requests padded.
140
+
141
+ Unlike ``pack_seq_triton`` this is dtype-agnostic (needed for the int32
142
+ slot/pos tensors) — it scatters with the shared per-step indices from
143
+ ``_build_decode_scatter_indices``. Used only for the non-uniform
144
+ (``requires_padding``) decode batch; uniform batches use a zero-copy
145
+ reshape.
146
+ """
147
+ req_id, intra = scatter_indices
148
+ out = torch.full(
149
+ (num_requests, lmax, *tokens.shape[1:]),
150
+ pad_value,
151
+ dtype=tokens.dtype,
152
+ device=tokens.device,
153
+ )
154
+ out[req_id, intra] = tokens
155
+ return out
156
+
157
+
158
+def _decode_topk_seq_lens(
159
+ positions: torch.Tensor,
160
+ decode_lens: torch.Tensor,
161
+ num_decode_tokens: int,
162
+ batch_size: int,
163
+ next_n: int,
164
+ requires_padding: bool,
165
+) -> torch.Tensor:
166
+ """Token-granular seq_len (pos + 1) per pool-topk row, layout-aware.
167
+
168
+ ``pool_topk`` (and the logits it comes from) follow the padded
169
+ ``[batch_size, next_n]`` grid whenever ``requires_padding`` is set, so row
170
+ ``(b, t)`` corresponds to flat decode token ``offset_b + t`` -- NOT
171
+ ``b * next_n + t``. Slicing flat ``positions[: batch_size * next_n]``
172
+ (the uniform-layout shortcut) misaligns every row after the first
173
+ non-uniform request and, past the decode region, reads prefill tokens'
174
+ positions; ``expand_pools_and_append_tail`` then anchors the tail at
175
+ another request's length, dropping the row's real tail tokens or emitting
176
+ indices past its sequence (out-of-bounds block-table reads). Padded rows
177
+ get 0 (empty tail); they are dropped by ``unpack_seq_triton`` anyway.
178
+ """
179
+ n = batch_size * next_n
180
+ if not requires_padding:
181
+ return positions[:n].to(torch.int32) + 1
182
+ scatter_idx = _build_decode_scatter_indices(
183
+ decode_lens, batch_size, num_decode_tokens
184
+ )
185
+ padded = _scatter_decode_tokens_by_request(
186
+ positions[:num_decode_tokens].to(torch.int32),
187
+ -1,
188
+ batch_size,
189
+ next_n,
190
+ scatter_idx,
191
+ )
192
+ return padded.reshape(n) + 1 # pad rows: -1 + 1 = 0 -> empty tail
193
+
194
+
195
+def _fill_causal_indices(rows: torch.Tensor, positions: torch.Tensor) -> None:
196
+ causal_range = torch.arange(rows.shape[1], device=rows.device, dtype=torch.int32)
197
+ positions = positions.to(torch.int32)
198
+ rows[:] = causal_range[None, :]
199
+ rows[causal_range[None, :] > positions[:, None]] = -1
200
+
201
+
202
+def _fill_short_decode_causal_indices(
203
+ topk_indices_buffer: torch.Tensor,
204
+ positions: torch.Tensor | None,
205
+ num_decode_tokens: int,
206
+ max_seq_len: int,
207
+ topk_tokens: int,
208
+) -> bool:
209
+ """Fill exact causal rows when sparse decode would select every token."""
210
+ if positions is None or positions.numel() == 0 or max_seq_len > topk_tokens:
211
+ return False
212
+ _fill_causal_indices(
213
+ topk_indices_buffer[:num_decode_tokens], positions[:num_decode_tokens]
214
+ )
215
+ return True
216
+
217
+
218
+def _gather_workspace_shapes(
219
+ total_seq_lens: int,
220
+ head_dim: int,
221
+ fp8_dtype: torch.dtype,
222
+ use_fp4_cache: bool,
223
+) -> tuple[tuple[tuple[int, int], torch.dtype], tuple[tuple[int, int], torch.dtype]]:
224
+ """Return ((values_shape, values_dtype), (scales_shape, scales_dtype)) for
225
+ the K-gather workspace. FP8 path: (T, head_dim) fp8 + (T, 4) uint8 fp32
226
+ scales. MXFP4 path: (T, head_dim // 2) uint8 packed mxfp4 +
227
+ (T, head_dim // MXFP4_BLOCK_SIZE) uint8 ue8m0 scales."""
228
+ if use_fp4_cache:
229
+ return (
230
+ ((total_seq_lens, head_dim // 2), torch.uint8),
231
+ ((total_seq_lens, head_dim // MXFP4_BLOCK_SIZE), torch.uint8),
232
+ )
233
+ return (
234
+ ((total_seq_lens, head_dim), fp8_dtype),
235
+ ((total_seq_lens, 4), torch.uint8),
236
+ )
237
+
238
+
239
+def kv_cache_as_quant_view(
240
+ kv_cache: torch.Tensor,
241
+ head_dim: int,
242
+ use_fp4_cache: bool,
243
+) -> torch.Tensor:
244
+ """4D ``[num_blocks, block_size, 1, head_width]`` view expected by
245
+ DeepGEMM, from the 3D indexer kv-cache allocation."""
246
+ if use_fp4_cache:
247
+ assert kv_cache.ndim == 3 and kv_cache.dtype == torch.uint8
248
+ num_blocks, block_size, _ = kv_cache.shape
249
+ page_bytes = int(kv_cache.stride(0))
250
+ fp4_bytes = head_dim // 2 + head_dim // MXFP4_BLOCK_SIZE
251
+ return torch.as_strided(
252
+ kv_cache,
253
+ size=(num_blocks, block_size, 1, fp4_bytes),
254
+ stride=(page_bytes, fp4_bytes, fp4_bytes, 1),
255
+ )
256
+ return kv_cache.unsqueeze(-2)
257
+
258
+
259
+@eager_break_during_capture
260
+def sparse_attn_indexer_kpool(
261
+ hidden_states: torch.Tensor,
262
+ k_cache_prefix: LayerNameType,
263
+ kv_cache: torch.Tensor,
264
+ q_quant: torch.Tensor,
265
+ q_scale: torch.Tensor | None,
266
+ k: torch.Tensor,
267
+ weights: torch.Tensor,
268
+ quant_block_size: int,
269
+ scale_fmt: str | None,
270
+ topk_tokens: int,
271
+ head_dim: int,
272
+ max_model_len: int,
273
+ total_seq_lens: int,
274
+ topk_indices_buffer: torch.Tensor,
275
+ skip_k_cache_insert: bool,
276
+ use_fp4_cache: bool = False,
277
+ # kpool params (Plan-A: gate is consumed at write time and read back at
278
+ # topk time to softmax-weight the pool).
279
+ gate_score: torch.Tensor | None = None,
280
+ compress_ape: torch.Tensor | None = None,
281
+ index_kpool: int = 1,
282
+ positions: torch.Tensor | None = None,
283
+ # Paged tail cache (in-progress pool's raw K + gate score), replacing the
284
+ # transient _DECODE_TAIL ring. tail_prefix resolves attn_metadata[tail_prefix]
285
+ # for the tail group's token-granular slot_mapping. None on the dummy/profiling
286
+ # path and when the tail cache is disabled.
287
+ tail_kv_cache: torch.Tensor | None = None,
288
+ tail_prefix: str | None = None,
289
+) -> torch.Tensor:
290
+ # careful! this will be None in dummy run
291
+ attn_metadata = get_forward_context().attn_metadata
292
+ fp8_dtype = current_platform.fp8_dtype()
293
+ k_cache_prefix = _resolve_layer_name(k_cache_prefix)
294
+
295
+ # assert isinstance(attn_metadata, dict)
296
+ if not isinstance(attn_metadata, dict):
297
+ # Reserve workspace for indexer during profiling run
298
+ values_spec, scales_spec = _gather_workspace_shapes(
299
+ total_seq_lens, head_dim, fp8_dtype, use_fp4_cache
300
+ )
301
+ current_workspace_manager().get_simultaneous(
302
+ values_spec,
303
+ scales_spec,
304
+ ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
305
+ )
306
+
307
+ # Reserve profiler-visible memory for the worst-case decode logits,
308
+ # whose shape is [B * next_n, max_model_len]. This profiling branch
309
+ # returns before invoking the logits kernel itself.
310
+ cfg = get_current_vllm_config_or_none()
311
+ worst_decode_tokens = 0
312
+ if cfg is not None:
313
+ sched = cfg.scheduler_config
314
+ num_spec = (
315
+ cfg.speculative_config.num_speculative_tokens
316
+ if cfg.speculative_config is not None
317
+ else 0
318
+ )
319
+ worst_decode_tokens = min(
320
+ sched.max_num_seqs * (num_spec + 1),
321
+ sched.max_num_batched_tokens,
322
+ )
323
+ # float32 logits -> 4 bytes/element; uint8 sentinel so elems == bytes.
324
+ decode_logits_elems = worst_decode_tokens * max_model_len * 4
325
+ prefill_cap_elems = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024
326
+ max_logits_elems = max(decode_logits_elems, prefill_cap_elems)
327
+ _ = torch.empty(
328
+ max_logits_elems, dtype=torch.uint8, device=hidden_states.device
329
+ )
330
+
331
+ return topk_indices_buffer
332
+ attn_metadata_narrowed = attn_metadata[k_cache_prefix]
333
+ assert isinstance(attn_metadata_narrowed, DeepseekV32IndexerMetadata)
334
+ slot_mapping = attn_metadata_narrowed.slot_mapping
335
+ has_decode = attn_metadata_narrowed.num_decodes > 0
336
+ has_prefill = attn_metadata_narrowed.num_prefills > 0
337
+ num_decode_tokens = attn_metadata_narrowed.num_decode_tokens
338
+
339
+ # q_scale is required iff the FP4 cache path is enabled; the FP8 path
340
+ # folds the Q scale into `weights` inside fused_indexer_q_rope_quant.
341
+ if use_fp4_cache:
342
+ assert q_scale is not None, "use_fp4_cache=True requires q_scale"
343
+ else:
344
+ assert q_scale is None, "q_scale must be None when use_fp4_cache=False"
345
+
346
+ # During speculative decoding, k may be padded to the CUDA graph batch
347
+ # size while slot_mapping only covers actual tokens. Truncate k to avoid
348
+ # out-of-bounds reads in the kernel.
349
+ num_tokens = slot_mapping.shape[0]
350
+ if k is not None:
351
+ k = k[:num_tokens]
352
+
353
+ if not skip_k_cache_insert:
354
+ assert not use_fp4_cache, "Unfused FP4 Insert is not supported yet"
355
+ if index_kpool > 1 and gate_score is not None and compress_ape is not None:
356
+ # kpool prefill write: pool kpool consecutive prefill tokens via
357
+ # softmax(gate+ape)-weighted sum -> Hadamard -> fp8 -> pool slots.
358
+ # Decode tokens (the first num_decode_tokens in the batch) cannot be
359
+ # pooled here — their pool's earlier tokens are not in this batch —
360
+ # so they are deferred to the tail-buffer kernel in has_decode.
361
+ # compress_ratio == index_kpool makes slot_mapping pool-granular.
362
+ n_prefill = num_tokens - num_decode_tokens
363
+ if n_prefill > 0:
364
+ # decode tokens are batched first; prefill tokens follow.
365
+ prefill_slice = slice(num_decode_tokens, num_tokens)
366
+ _kpool_compress_insert(
367
+ k[prefill_slice],
368
+ gate_score[prefill_slice],
369
+ compress_ape,
370
+ kv_cache,
371
+ slot_mapping[prefill_slice],
372
+ index_kpool,
373
+ head_dim,
374
+ round_scale=(scale_fmt is not None),
375
+ )
376
+ # Persist each request's incomplete prefill pool so decode can
377
+ # finish it, including after PD transfer. Tail slots use
378
+ # ``pos % kpool`` within the request's tail block. Processing
379
+ # only the batch's trailing tokens would miss all but the last
380
+ # request in a multi-request prefill.
381
+ if tail_kv_cache is not None and tail_prefix is not None:
382
+ tail_meta = attn_metadata.get(_resolve_layer_name(tail_prefix))
383
+ if tail_meta is not None:
384
+ assert isinstance(tail_meta, DeepseekV32IndexerMetadata)
385
+ kpool_ops.kpool_seed_tail_cache(
386
+ tail_kv_cache,
387
+ k[prefill_slice],
388
+ gate_score[prefill_slice],
389
+ tail_meta.slot_mapping[prefill_slice],
390
+ index_kpool,
391
+ head_dim,
392
+ )
393
+ else:
394
+ # standard: per-token fp8 quant + scatter (all tokens).
395
+ assert scale_fmt is not None
396
+ if current_platform.is_rocm():
397
+ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
398
+ indexer_k_quant_and_cache_triton,
399
+ )
400
+
401
+ indexer_k_quant_and_cache_triton(
402
+ k,
403
+ kv_cache,
404
+ slot_mapping,
405
+ quant_block_size,
406
+ scale_fmt,
407
+ )
408
+ else:
409
+ ops.indexer_k_quant_and_cache(
410
+ k,
411
+ kv_cache,
412
+ slot_mapping,
413
+ quant_block_size,
414
+ scale_fmt,
415
+ )
416
+
417
+ topk_indices_buffer[: hidden_states.shape[0]] = -1
418
+ if has_prefill:
419
+ prefill_metadata = attn_metadata_narrowed.prefill
420
+ assert prefill_metadata is not None
421
+
422
+ # Short sequences select every pool, so skip sparse scoring and fill
423
+ # the top-k buffer with all causal token indices. The index-K cache was
424
+ # already written above.
425
+ n_prefill_sf = num_tokens - num_decode_tokens
426
+ # Host-side short-prefill predicate: max_prefill_seq_len is computed
427
+ # in the metadata builder (exact for prefill rows) and equals
428
+ # positions[prefill_slice].max() + 1, so this replaces a
429
+ # positions.max().item() device sync per layer. -1 (unknown metadata)
430
+ # falls back to the device-side check.
431
+ if prefill_metadata.max_prefill_seq_len >= 0:
432
+ short_prefill = (
433
+ n_prefill_sf > 0
434
+ and positions is not None
435
+ and prefill_metadata.max_prefill_seq_len <= topk_tokens
436
+ )
437
+ else:
438
+ short_prefill = (
439
+ n_prefill_sf > 0
440
+ and positions is not None
441
+ and int(positions[num_decode_tokens:num_tokens].max().item()) + 1
442
+ <= topk_tokens
443
+ )
444
+ if short_prefill:
445
+ # short_prefill is only True when positions is not None (above),
446
+ # but narrow explicitly for the indexer below.
447
+ assert positions is not None
448
+ _pos = positions[num_decode_tokens:num_tokens].to(torch.int32)
449
+ _buf = topk_indices_buffer[num_decode_tokens:num_tokens]
450
+ _fill_causal_indices(_buf, _pos)
451
+
452
+ # Get the full shared workspace buffers once (will allocate on first use).
453
+ # Layout switches between FP8 (head_dim bytes + 4-byte fp32 scale) and
454
+ # MXFP4 (head_dim/2 bytes packed + head_dim/MXFP4_BLOCK_SIZE ue8m0
455
+ # scales) based on use_fp4_cache.
456
+ workspace_manager = current_workspace_manager()
457
+ values_spec, scales_spec = _gather_workspace_shapes(
458
+ total_seq_lens, head_dim, fp8_dtype, use_fp4_cache
459
+ )
460
+ k_quant_full, k_scale_full = workspace_manager.get_simultaneous(
461
+ values_spec,
462
+ scales_spec,
463
+ )
464
+ for chunk in prefill_metadata.chunks if not short_prefill else ():
465
+ k_quant = k_quant_full[: chunk.total_seq_lens]
466
+ k_scale = k_scale_full[: chunk.total_seq_lens]
467
+
468
+ if not chunk.skip_kv_gather:
469
+ if current_platform.is_rocm():
470
+ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
471
+ cp_gather_indexer_k_quant_cache_triton,
472
+ )
473
+
474
+ cp_gather_indexer_k_quant_cache_triton(
475
+ kv_cache,
476
+ k_quant,
477
+ k_scale,
478
+ chunk.block_table,
479
+ chunk.cu_seq_lens,
480
+ token_to_seq=chunk.token_to_seq,
481
+ )
482
+ else:
483
+ ops.cp_gather_indexer_k_quant_cache(
484
+ kv_cache,
485
+ k_quant,
486
+ k_scale,
487
+ chunk.block_table,
488
+ chunk.cu_seq_lens,
489
+ )
490
+
491
+ q_slice = q_quant[chunk.token_start : chunk.token_end]
492
+ q_scale_slice = (
493
+ q_scale[chunk.token_start : chunk.token_end]
494
+ if q_scale is not None
495
+ else None
496
+ )
497
+ # DeepGEMM scalar-type tags (zero-copy): MXFP4 values → int8
498
+ # (kPackedFP4), scales → int32 squeezed to 1-D kv_sf / 2-D q_sf.
499
+ if use_fp4_cache:
500
+ q_slice_cast = q_slice.view(torch.int8)
501
+ k_quant_cast = k_quant.view(torch.int8)
502
+ k_scale_cast = k_scale.view(torch.int32).squeeze(-1)
503
+ else:
504
+ q_slice_cast = q_slice
505
+ k_quant_cast = k_quant
506
+ k_scale_cast = k_scale.view(torch.float32).squeeze(-1)
507
+ if current_platform.is_rocm():
508
+ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
509
+ rocm_fp8_mqa_logits,
510
+ )
511
+
512
+ assert q_scale_slice is None
513
+ logits = rocm_fp8_mqa_logits(
514
+ q_slice_cast,
515
+ (k_quant_cast, k_scale_cast),
516
+ weights[chunk.token_start : chunk.token_end],
517
+ chunk.cu_seqlen_ks,
518
+ chunk.cu_seqlen_ke,
519
+ )
520
+ else:
521
+ from vllm.utils.deep_gemm import fp8_fp4_mqa_logits
522
+
523
+ logits = fp8_fp4_mqa_logits(
524
+ (q_slice_cast, q_scale_slice),
525
+ (k_quant_cast, k_scale_cast),
526
+ weights[chunk.token_start : chunk.token_end],
527
+ chunk.cu_seqlen_ks,
528
+ chunk.cu_seqlen_ke,
529
+ clean_logits=False,
530
+ )
531
+ num_rows = logits.shape[0]
532
+
533
+ # kpool: logits are pool-granular (compress_ratio == index_kpool),
534
+ # so topk selects pools. We pick topk_tokens // kpool pools then
535
+ # expand each pool back to its kpool constituent tokens.
536
+ select_k = topk_tokens // index_kpool if index_kpool > 1 else topk_tokens
537
+ if index_kpool > 1:
538
+ pool_topk = torch.empty(
539
+ (num_rows, select_k), dtype=torch.int32, device=logits.device
540
+ )
541
+ topk_dst = pool_topk
542
+ else:
543
+ topk_dst = topk_indices_buffer[
544
+ chunk.token_start : chunk.token_end, :topk_tokens
545
+ ]
546
+
547
+ if current_platform.is_xpu():
548
+ xpu_ops.top_k_per_row_prefill( # type: ignore[attr-defined]
549
+ logits,
550
+ chunk.cu_seqlen_ks,
551
+ chunk.cu_seqlen_ke,
552
+ topk_dst,
553
+ num_rows,
554
+ logits.stride(0),
555
+ logits.stride(1),
556
+ select_k,
557
+ )
558
+ else:
559
+ torch.ops._C.top_k_per_row_prefill(
560
+ logits,
561
+ chunk.cu_seqlen_ks,
562
+ chunk.cu_seqlen_ke,
563
+ topk_dst,
564
+ num_rows,
565
+ logits.stride(0),
566
+ logits.stride(1),
567
+ select_k,
568
+ )
569
+
570
+ if index_kpool > 1:
571
+ pool_ids = pool_topk.to(torch.int64)
572
+ if positions is not None:
573
+ # Fused expand-pools + append-tail into one Triton kernel
574
+ # (replaces ~25 elementwise ops). seq_len is token-granular
575
+ # (pos+1); the kernel derives pool_len internally.
576
+ q_seq = (
577
+ positions[chunk.token_start : chunk.token_end].to(torch.int32)
578
+ + 1
579
+ )
580
+ expanded = kpool_ops.expand_pools_and_append_tail(
581
+ pool_ids, q_seq, index_kpool
582
+ )
583
+ else:
584
+ valid = pool_ids >= 0
585
+ expanded = kpool_ops.expand_pools_to_tokens(
586
+ pool_ids, valid, topk_tokens, index_kpool
587
+ )
588
+ topk_indices_buffer[
589
+ chunk.token_start : chunk.token_end, : expanded.shape[-1]
590
+ ] = expanded
591
+
592
+ if has_decode:
593
+ decode_metadata = attn_metadata_narrowed.decode
594
+ assert decode_metadata is not None
595
+ kv_cache_raw = kv_cache # raw [num_blocks, block_size, head_dim+4] for writes
596
+ kv_cache = kv_cache_as_quant_view(kv_cache, head_dim, use_fp4_cache)
597
+
598
+ # Update the tail before reading logits; completed pools are compressed
599
+ # into the slot supplied by slot_mapping.
600
+ # Spec verification groups tokens by request and preserves position
601
+ # order so each token is stashed before the next completes its pool.
602
+ # Positions must remain token-granular because the kernel derives the
603
+ # pool phase and tail index from ``pos % kpool``.
604
+ if (
605
+ index_kpool > 1
606
+ and gate_score is not None
607
+ and compress_ape is not None
608
+ and positions is not None
609
+ and not skip_k_cache_insert
610
+ ):
611
+ num_requests = attn_metadata_narrowed.num_decodes
612
+ # Kpool writes must recover the original request grouping after the
613
+ # indexer's flattened decode path. Host metadata avoids a CUDA graph
614
+ # sync when choosing the uniform or padded layout.
615
+ per_req_lens = decode_metadata.per_req_decode_lens
616
+ if per_req_lens is not None:
617
+ use_uniform = (
618
+ decode_metadata.decode_is_uniform
619
+ and num_decode_tokens
620
+ == num_requests * decode_metadata.write_max_decode_len
621
+ )
622
+ group_lens = per_req_lens
623
+ lmax = decode_metadata.write_max_decode_len
624
+ else:
625
+ # Legacy metadata without per-request lens: fall back to the
626
+ # host-side requires_padding flag. Unreached now (per-request
627
+ # lens is always populated for decode), kept defensive.
628
+ use_uniform = not decode_metadata.requires_padding
629
+ group_lens = decode_metadata.decode_lens
630
+ lmax = int(decode_metadata.decode_lens.max().item())
631
+ if not use_uniform:
632
+ # Non-uniform decode_lens (mixed plain-decode + spec-verify, or
633
+ # a variable MTP-verify batch): scatter actual tokens into a
634
+ # padded [B, lmax] layout. int32 tensors can't go through
635
+ # pack_seq_triton (float/uint8 only). The scatter indices are
636
+ # shared by all five scatters below (and the tail slot one).
637
+ scatter_idx = _build_decode_scatter_indices(
638
+ group_lens, num_requests, num_decode_tokens
639
+ )
640
+ dec_k = _scatter_decode_tokens_by_request(
641
+ k[:num_decode_tokens], 0, num_requests, lmax, scatter_idx
642
+ )
643
+ dec_gate = _scatter_decode_tokens_by_request(
644
+ gate_score[:num_decode_tokens],
645
+ 0,
646
+ num_requests,
647
+ lmax,
648
+ scatter_idx,
649
+ )
650
+ dec_slot = _scatter_decode_tokens_by_request(
651
+ slot_mapping[:num_decode_tokens],
652
+ -1,
653
+ num_requests,
654
+ lmax,
655
+ scatter_idx,
656
+ )
657
+ dec_pos = _scatter_decode_tokens_by_request(
658
+ positions[:num_decode_tokens].to(torch.int32),
659
+ -1,
660
+ num_requests,
661
+ lmax,
662
+ scatter_idx,
663
+ )
664
+ else:
665
+ next_n = num_decode_tokens // num_requests
666
+ shape2 = (num_requests, next_n)
667
+ dec_k = k[:num_decode_tokens].view(*shape2, head_dim)
668
+ dec_gate = gate_score[:num_decode_tokens].view(*shape2, head_dim)
669
+ dec_slot = slot_mapping[:num_decode_tokens].view(shape2)
670
+ dec_pos = positions[:num_decode_tokens].to(torch.int32).view(shape2)
671
+ tail_meta = (
672
+ attn_metadata.get(_resolve_layer_name(tail_prefix))
673
+ if tail_prefix is not None
674
+ else None
675
+ )
676
+ # Paged tail cache replaces the transient _DECODE_TAIL ring. Group
677
+ # the tail group's token-granular slot_mapping per-request, mirroring
678
+ # dec_slot / dec_pos, so the kernel gets each request's current-token
679
+ # tail slot (block * kpool + pos % kpool).
680
+ if tail_meta is not None:
681
+ assert isinstance(tail_meta, DeepseekV32IndexerMetadata)
682
+ if tail_meta is None or tail_kv_cache is None:
683
+ dec_tail_slot = None
684
+ elif not use_uniform:
685
+ dec_tail_slot = _scatter_decode_tokens_by_request(
686
+ tail_meta.slot_mapping[:num_decode_tokens],
687
+ -1,
688
+ num_requests,
689
+ lmax,
690
+ scatter_idx,
691
+ )
692
+ else:
693
+ dec_tail_slot = tail_meta.slot_mapping[:num_decode_tokens].view(shape2)
694
+ # The compress kernel writes the raw fp8 cache (not the quant view);
695
+ # pass the underlying kv_cache, not kv_cache_quant_view.
696
+ if dec_tail_slot is not None:
697
+ # Single batched launch over [num_requests, next_n] replaces the
698
+ # per-token sequential loop. The kernel iterates each request's
699
+ # tokens in position order internally, preserving the
700
+ # pool-completion read-after-stash dependency that the loop
701
+ # provided. Inputs are already grouped per request (uniform:
702
+ # view; non-uniform: _scatter_decode_tokens_by_request padded to
703
+ # [B, lmax]) — no per-token .contiguous() copies needed.
704
+ kpool_ops.kpool_decode_update_and_maybe_write_cache_batched(
705
+ kv_cache_raw,
706
+ tail_kv_cache,
707
+ dec_tail_slot,
708
+ dec_k,
709
+ dec_gate,
710
+ compress_ape,
711
+ dec_slot,
712
+ dec_pos,
713
+ index_kpool,
714
+ head_dim,
715
+ round_scale=(scale_fmt is not None),
716
+ )
717
+ if current_platform.is_cuda_alike() and _fill_short_decode_causal_indices(
718
+ topk_indices_buffer,
719
+ positions,
720
+ num_decode_tokens,
721
+ attn_metadata_narrowed.max_seq_len,
722
+ topk_tokens,
723
+ ):
724
+ return topk_indices_buffer
725
+ decode_lens = decode_metadata.decode_lens
726
+ if decode_metadata.requires_padding:
727
+ # Padding also covers short chunked prefills classified as decode.
728
+ # MXFP4 uses zero-byte padding so padded slots dequantize to zero.
729
+ if q_scale is not None:
730
+ padded_q_quant_decode_tokens = pack_seq_triton(
731
+ q_quant[:num_decode_tokens], decode_lens, pad_value=0
732
+ )
733
+ padded_q_scale = pack_seq_triton(
734
+ q_scale[:num_decode_tokens], decode_lens, pad_value=0
735
+ )
736
+ else:
737
+ padded_q_quant_decode_tokens = pack_seq_triton(
738
+ q_quant[:num_decode_tokens], decode_lens
739
+ )
740
+ padded_q_scale = None
741
+ padded_weights = pack_seq_triton(
742
+ weights[:num_decode_tokens], decode_lens, pad_value=0
743
+ ).reshape(-1, *weights.shape[1:])
744
+ else:
745
+ padded_q_quant_decode_tokens = q_quant[:num_decode_tokens].reshape(
746
+ decode_lens.shape[0], -1, *q_quant.shape[1:]
747
+ )
748
+ if q_scale is not None:
749
+ padded_q_scale = q_scale[:num_decode_tokens].reshape(
750
+ decode_lens.shape[0], -1, *q_scale.shape[1:]
751
+ )
752
+ else:
753
+ padded_q_scale = None
754
+ padded_weights = weights[:num_decode_tokens]
755
+ # TODO: move and optimize below logic with triton kernels
756
+ batch_size = padded_q_quant_decode_tokens.shape[0]
757
+ next_n = padded_q_quant_decode_tokens.shape[1]
758
+ num_padded_tokens = batch_size * next_n
759
+ seq_lens = decode_metadata.seq_lens[:batch_size]
760
+ # seq_lens is always 2D: (B, next_n) for native spec decode, (B, 1)
761
+ # otherwise. deep_gemm fp8_fp4_paged_mqa_logits requires 2D context_lens;
762
+ # the downstream topk kernels accept both 1D and 2D.
763
+ padded_q_quant_cast = (
764
+ padded_q_quant_decode_tokens.view(torch.int8)
765
+ if use_fp4_cache
766
+ else padded_q_quant_decode_tokens
767
+ )
768
+ if current_platform.is_rocm():
769
+ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
770
+ rocm_fp8_paged_mqa_logits,
771
+ )
772
+
773
+ assert padded_q_scale is None
774
+ logits = rocm_fp8_paged_mqa_logits(
775
+ padded_q_quant_cast,
776
+ kv_cache,
777
+ padded_weights[:num_padded_tokens],
778
+ seq_lens,
779
+ decode_metadata.block_table,
780
+ decode_metadata.schedule_metadata,
781
+ max_model_len=max_model_len,
782
+ )
783
+ else:
784
+ from vllm.utils.deep_gemm import fp8_fp4_paged_mqa_logits
785
+
786
+ logits = fp8_fp4_paged_mqa_logits(
787
+ (padded_q_quant_cast, padded_q_scale),
788
+ kv_cache,
789
+ padded_weights[:num_padded_tokens],
790
+ seq_lens,
791
+ decode_metadata.block_table,
792
+ decode_metadata.schedule_metadata,
793
+ max_model_len=max_model_len,
794
+ clean_logits=False,
795
+ )
796
+ num_rows = logits.shape[0]
797
+ # kpool: logits are pool-granular -> select topk_tokens//kpool pools,
798
+ # then expand each pool back to its kpool tokens.
799
+ select_k = topk_tokens // index_kpool if index_kpool > 1 else topk_tokens
800
+ if index_kpool > 1:
801
+ pool_topk = torch.empty(
802
+ (num_rows, select_k), dtype=torch.int32, device=logits.device
803
+ )
804
+ topk_dst = pool_topk
805
+ else:
806
+ topk_dst = topk_indices_buffer[:num_padded_tokens, :topk_tokens]
807
+
808
+ if current_platform.is_cuda() and select_k in (512, 1024, 2048):
809
+ workspace_manager = current_workspace_manager()
810
+ (topk_workspace,) = workspace_manager.get_simultaneous(
811
+ ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
812
+ )
813
+ torch.ops._C.persistent_topk(
814
+ logits,
815
+ seq_lens,
816
+ topk_dst,
817
+ topk_workspace,
818
+ select_k,
819
+ attn_metadata_narrowed.max_seq_len,
820
+ )
821
+ else:
822
+ if current_platform.is_xpu():
823
+ xpu_ops.top_k_per_row_decode( # type: ignore[attr-defined]
824
+ logits,
825
+ next_n,
826
+ seq_lens,
827
+ topk_dst,
828
+ num_rows,
829
+ logits.stride(0),
830
+ logits.stride(1),
831
+ select_k,
832
+ )
833
+ else:
834
+ torch.ops._C.top_k_per_row_decode(
835
+ logits,
836
+ next_n,
837
+ seq_lens,
838
+ topk_dst,
839
+ num_rows,
840
+ logits.stride(0),
841
+ logits.stride(1),
842
+ select_k,
843
+ )
844
+
845
+ # Resolve to token-level indices in the output buffer.
846
+ if index_kpool > 1:
847
+ pool_ids = pool_topk.to(torch.int64)
848
+ n = pool_topk.shape[0]
849
+ # Decode seq_lens are pool-granular; recover token lengths from
850
+ # positions using the padded [B, next_n] row layout when needed.
851
+ if positions is not None:
852
+ dec_seq = _decode_topk_seq_lens(
853
+ positions,
854
+ decode_lens,
855
+ num_decode_tokens,
856
+ batch_size,
857
+ next_n,
858
+ decode_metadata.requires_padding,
859
+ )
860
+ else:
861
+ dec_seq = decode_metadata.seq_lens[:n]
862
+ if dec_seq.ndim == 2:
863
+ dec_seq = dec_seq[:, -1]
864
+ dec_seq = dec_seq.to(torch.int32)
865
+ out = kpool_ops.expand_pools_and_append_tail(pool_ids, dec_seq, index_kpool)
866
+ else:
867
+ out = topk_dst
868
+
869
+ if decode_metadata.requires_padding:
870
+ # Drop padded query rows introduced by the next_n padding above.
871
+ out = unpack_seq_triton(
872
+ out.reshape(batch_size, -1, out.shape[-1]), decode_lens
873
+ )
874
+ topk_indices_buffer[: out.shape[0], : out.shape[-1]] = out
875
+
876
+ return topk_indices_buffer
877
+
878
+
879
+@CustomOp.register("sparse_attn_indexer_kpool")
880
+class SparseAttnIndexerKpool(CustomOp):
881
+ """Sparse Attention Indexer Custom Op Layer. This layer is extracted as a
882
+ separate custom op since it involves heavy custom kernels like `mqa_logits`,
883
+ `paged_mqa_logits` and `top_k_per_row`, etc. Those kernels maybe requires
884
+ specific memory layout or implementation for different hardware backends to
885
+ achieve optimal performance.
886
+
887
+ For now, the default native path will use CUDA backend path. Other platform
888
+ may requires add the corresponding Custom Op name `sparse_attn_indexer` to
889
+ `custom_ops` in `CompilationConfig` to enable the platform specific path.
890
+ """
891
+
892
+ def __init__(
893
+ self,
894
+ k_cache,
895
+ quant_block_size: int,
896
+ scale_fmt: str,
897
+ topk_tokens: int,
898
+ head_dim: int,
899
+ max_model_len: int,
900
+ max_total_seq_len: int,
901
+ topk_indices_buffer: torch.Tensor,
902
+ skip_k_cache_insert: bool = False,
903
+ use_fp4_cache: bool = False,
904
+ tail_cache=None,
905
+ ):
906
+ super().__init__()
907
+ self.k_cache = k_cache
908
+ self.tail_cache = tail_cache
909
+ self.quant_block_size = quant_block_size
910
+ self.scale_fmt = scale_fmt
911
+ self.topk_tokens = topk_tokens
912
+ self.head_dim = head_dim
913
+ self.max_model_len = max_model_len
914
+ self.max_total_seq_len = max_total_seq_len
915
+ self.topk_indices_buffer = topk_indices_buffer
916
+ self.skip_k_cache_insert = skip_k_cache_insert
917
+ self.use_fp4_cache = use_fp4_cache
918
+ if current_platform.is_cuda() and not has_deep_gemm():
919
+ raise RuntimeError(
920
+ "Sparse Attention Indexer CUDA op requires DeepGEMM to be installed."
921
+ )
922
+
923
+ def forward_native(
924
+ self,
925
+ hidden_states: torch.Tensor,
926
+ q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
927
+ k: torch.Tensor,
928
+ weights: torch.Tensor,
929
+ *,
930
+ gate_score: torch.Tensor | None = None,
931
+ compress_ape: torch.Tensor | None = None,
932
+ index_kpool: int = 1,
933
+ positions: torch.Tensor | None = None,
934
+ ):
935
+ if current_platform.is_cuda() or current_platform.is_xpu():
936
+ return self.forward_cuda(
937
+ hidden_states,
938
+ q_quant,
939
+ k,
940
+ weights,
941
+ gate_score=gate_score,
942
+ compress_ape=compress_ape,
943
+ index_kpool=index_kpool,
944
+ positions=positions,
945
+ )
946
+ elif current_platform.is_rocm():
947
+ return self.forward_hip(
948
+ hidden_states,
949
+ q_quant,
950
+ k,
951
+ weights,
952
+ gate_score=gate_score,
953
+ compress_ape=compress_ape,
954
+ index_kpool=index_kpool,
955
+ positions=positions,
956
+ )
957
+ else:
958
+ raise NotImplementedError(
959
+ "SparseAttnIndexer native forward is only implemented for "
960
+ "CUDA, ROCm and XPU platforms."
961
+ )
962
+
963
+ def forward_cuda(
964
+ self,
965
+ hidden_states: torch.Tensor,
966
+ q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
967
+ k: torch.Tensor,
968
+ weights: torch.Tensor,
969
+ *,
970
+ gate_score: torch.Tensor | None = None,
971
+ compress_ape: torch.Tensor | None = None,
972
+ index_kpool: int = 1,
973
+ positions: torch.Tensor | None = None,
974
+ ):
975
+ # FP8 path: single tensor (per-token scale is folded into `weights`).
976
+ # FP4 path: (values, scales) tuple with scales required by the kernel.
977
+ if isinstance(q_quant, tuple):
978
+ q_values, q_scale = q_quant
979
+ else:
980
+ q_values, q_scale = q_quant, None
981
+ return sparse_attn_indexer_kpool(
982
+ hidden_states,
983
+ self.k_cache.prefix,
984
+ self.k_cache.kv_cache,
985
+ q_values,
986
+ q_scale,
987
+ k,
988
+ weights,
989
+ self.quant_block_size,
990
+ self.scale_fmt,
991
+ self.topk_tokens,
992
+ self.head_dim,
993
+ self.max_model_len,
994
+ self.max_total_seq_len,
995
+ self.topk_indices_buffer,
996
+ self.skip_k_cache_insert,
997
+ self.use_fp4_cache,
998
+ gate_score,
999
+ compress_ape,
1000
+ index_kpool,
1001
+ positions,
1002
+ self.tail_cache.kv_cache if self.tail_cache is not None else None,
1003
+ self.tail_cache.prefix if self.tail_cache is not None else None,
1004
+ )
1005
+
1006
+ def forward_hip(
1007
+ self,
1008
+ hidden_states: torch.Tensor,
1009
+ q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
1010
+ k: torch.Tensor,
1011
+ weights: torch.Tensor,
1012
+ *,
1013
+ gate_score: torch.Tensor | None = None,
1014
+ compress_ape: torch.Tensor | None = None,
1015
+ index_kpool: int = 1,
1016
+ positions: torch.Tensor | None = None,
1017
+ ):
1018
+ assert not self.use_fp4_cache, "AMD platform doesn't support fp4 cache yet"
1019
+ assert isinstance(q_quant, torch.Tensor), (
1020
+ "AMD sparse_attn_indexer expects a single FP8 q_quant tensor"
1021
+ )
1022
+ if rocm_aiter_ops.is_enabled():
1023
+ if index_kpool <= 1:
1024
+ return torch.ops.vllm.rocm_aiter_sparse_attn_indexer(
1025
+ hidden_states,
1026
+ _encode_layer_name(self.k_cache.prefix),
1027
+ self.k_cache.kv_cache,
1028
+ q_quant,
1029
+ k,
1030
+ weights,
1031
+ self.quant_block_size,
1032
+ self.scale_fmt,
1033
+ self.topk_tokens,
1034
+ self.head_dim,
1035
+ self.max_model_len,
1036
+ self.max_total_seq_len,
1037
+ self.topk_indices_buffer,
1038
+ skip_k_cache_insert=self.skip_k_cache_insert,
1039
+ )
1040
+ return self.forward_cuda(
1041
+ hidden_states,
1042
+ q_quant,
1043
+ k,
1044
+ weights,
1045
+ gate_score=gate_score,
1046
+ compress_ape=compress_ape,
1047
+ index_kpool=index_kpool,
1048
+ positions=positions,
1049
+ )
1050
+ raise RuntimeError(
1051
+ "Sparse attention indexer ROCm path is only supported on AITER. "
1052
+ "Please enable aiter with VLLM_ROCM_USE_AITER=1"
1053
+ )
vllm/model_executor/models/registry.py
+6
@@ -120,6 +120,7 @@ _TEXT_GENERATION_MODELS = {
120
"Glm4MoeForCausalLM": ("glm4_moe", "Glm4MoeForCausalLM"),
121
"Glm4MoeLiteForCausalLM": ("glm4_moe_lite", "Glm4MoeLiteForCausalLM"),
122
"GlmMoeDsaForCausalLM": ("vllm.models.deepseek_v32", "GlmMoeDsaForCausalLM"),
123
+ "Glm5NextForCausalLM": ("vllm.models.glm5next", "Glm5NextForCausalLM"),
124
"GptOssForCausalLM": ("gpt_oss", "GptOssForCausalLM"),
125
"GPT2LMHeadModel": ("gpt2", "GPT2LMHeadModel"),
126
"GPTJForCausalLM": ("gpt_j", "GPTJForCausalLM"),
@@ -413,6 +414,10 @@ _MULTIMODAL_MODELS = {
414
"Glm4vForConditionalGeneration": ("glm4_1v", "Glm4vForConditionalGeneration"),
415
"Glm4vMoeForConditionalGeneration": ("glm4_1v", "Glm4vMoeForConditionalGeneration"),
416
"GlmOcrForConditionalGeneration": ("glm_ocr", "GlmOcrForConditionalGeneration"),
417
+ "Glm5NextForConditionalGeneration": (
418
+ "vllm.models.glm5next",
419
+ "Glm5NextForConditionalGeneration",
420
+ ),
421
"GraniteSpeechForConditionalGeneration": (
422
"granite_speech",
423
"GraniteSpeechForConditionalGeneration",
@@ -666,6 +671,7 @@ _SPECULATIVE_DECODING_MODELS = {
671
"Glm4MoeMTPModel": ("glm4_moe_mtp", "Glm4MoeMTP"),
672
"Glm4MoeLiteMTPModel": ("glm4_moe_lite_mtp", "Glm4MoeLiteMTP"),
673
"GlmOcrMTPModel": ("glm_ocr_mtp", "GlmOcrMTP"),
674
+ "Glm5NextMTPModel": ("vllm.models.glm5next", "Glm5NextMTP"),
675
"MedusaModel": ("medusa", "Medusa"),
676
"OpenPanguMTPModel": ("openpangu_mtp", "OpenPanguMTP"),
677
"Qwen3NextMTP": ("qwen3_next_mtp", "Qwen3NextMTP"),
vllm/model_executor/warmup/kernel_warmup.py
+4
@@ -41,6 +41,9 @@ from vllm.model_executor.warmup.qwen_triton_warmup import qwen_triton_warmup
41
from vllm.model_executor.warmup.replayssm_warmup import (
42
replayssm_autotune_warmup,
43
)
44
+from vllm.model_executor.warmup.spec_decode_rejection_warmup import (
45
+ spec_decode_rejection_warmup,
46
+)
47
from vllm.platforms import current_platform
48
from vllm.utils.deep_gemm import is_deep_gemm_supported
49
from vllm.utils.flashinfer import has_flashinfer
@@ -163,6 +166,7 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):
166
if worker.vllm_config.kernel_config.enable_jit_warmup:
167
kimi_k3_triton_warmup(worker)
168
fa4_cutedsl_warmup(worker)
169
+ spec_decode_rejection_warmup(worker)
170
qwen4_exp_qsa_triton_warmup(worker)
171
172
if current_platform.has_device_capability(90):
vllm/model_executor/warmup/spec_decode_rejection_warmup.py
new
+111
@@ -0,0 +1,111 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Warm up spec-decode rejection-sampler Triton kernels.
4
+
5
+The rejection sampler kernels (``_compute_local_logits_stats_kernel``,
6
+``_rejection_kernel``, ``_resample_kernel``) are JIT-compiled by Triton on
7
+first use. Without warmup, the first spec-decode request pays a multi-second
8
+compilation cost. This pre-compiles them with dummy data matching the
9
+server's vocab size and speculative config.
10
+"""
11
+
12
+from __future__ import annotations
13
+
14
+from typing import TYPE_CHECKING
15
+
16
+import torch
17
+
18
+from vllm.logger import init_logger
19
+
20
+if TYPE_CHECKING:
21
+ from vllm.v1.worker.gpu_worker import Worker
22
+
23
+logger = init_logger(__name__)
24
+
25
+
26
+@torch.inference_mode()
27
+def spec_decode_rejection_warmup(worker: Worker) -> None:
28
+ spec_config = worker.vllm_config.speculative_config
29
+ if spec_config is None:
30
+ return
31
+
32
+ from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import (
33
+ rejection_sample,
34
+ )
35
+
36
+ model_config = worker.vllm_config.model_config
37
+ vocab_size = model_config.get_vocab_size()
38
+ num_spec = spec_config.num_speculative_tokens
39
+ if num_spec <= 0 or vocab_size <= 0:
40
+ return
41
+
42
+ # Mirror the constexpr-relevant flags the runtime uses.
43
+ rejection_method = getattr(spec_config, "rejection_sample_method", None)
44
+ use_block_verification = rejection_method == "block"
45
+ use_synthetic = rejection_method == "synthetic"
46
+
47
+ device = torch.device("cuda")
48
+ num_reqs = 1
49
+ tokens_per_req = num_spec + 1
50
+ num_logits = num_reqs * tokens_per_req
51
+
52
+ # Triton JIT-specializes on tensor dtypes. The target logits may be fp32
53
+ # (apply_sampling_params copies to fp32 when processing is needed) or the
54
+ # model dtype (pass-through otherwise), while draft logits are always the
55
+ # model dtype. Warm every (target, draft) combination the runtime can hit.
56
+ model_dtype = model_config.dtype
57
+ warmup_dtype_pairs = {
58
+ (model_dtype, model_dtype),
59
+ (torch.float32, torch.float32),
60
+ (torch.float32, model_dtype),
61
+ (model_dtype, torch.float32),
62
+ }
63
+
64
+ logger.info(
65
+ "Warming up spec-decode rejection sampler kernels "
66
+ "(vocab=%d, num_spec=%d, dtype_pairs=%s, block_verify=%s).",
67
+ vocab_size,
68
+ num_spec,
69
+ [(str(t), str(d)) for t, d in warmup_dtype_pairs],
70
+ use_block_verification,
71
+ )
72
+ for tgt_dtype, draft_dtype in warmup_dtype_pairs:
73
+ target_logits = torch.zeros(
74
+ (num_logits, vocab_size), dtype=tgt_dtype, device=device
75
+ )
76
+ draft_logits = torch.zeros(
77
+ (num_reqs, num_spec, vocab_size), dtype=draft_dtype, device=device
78
+ )
79
+ synthetic_rates = (
80
+ torch.full((num_spec,), 0.5, dtype=torch.float32, device=device)
81
+ if use_synthetic
82
+ else None
83
+ )
84
+ try:
85
+ rejection_sample(
86
+ target_logits=target_logits,
87
+ draft_logits=draft_logits,
88
+ draft_sampled=torch.zeros(num_logits, dtype=torch.int64, device=device),
89
+ cu_num_logits=torch.tensor(
90
+ [0, num_logits], dtype=torch.int32, device=device
91
+ ),
92
+ pos=torch.zeros(num_logits, dtype=torch.int64, device=device),
93
+ idx_mapping=torch.zeros(num_reqs, dtype=torch.int32, device=device),
94
+ expanded_idx_mapping=torch.zeros(
95
+ num_logits, dtype=torch.int32, device=device
96
+ ),
97
+ expanded_local_pos=torch.arange(
98
+ num_logits, dtype=torch.int32, device=device
99
+ ),
100
+ temperature=torch.zeros(num_reqs, dtype=torch.float32, device=device),
101
+ seed=torch.full((num_reqs,), 42, dtype=torch.int64, device=device),
102
+ num_speculative_steps=num_spec,
103
+ synthetic_conditional_rates=synthetic_rates,
104
+ use_fp64=False,
105
+ use_block_verification=use_block_verification,
106
+ )
107
+ except Exception:
108
+ logger.warning(
109
+ "Skipping spec-decode rejection sampler warmup.", exc_info=True
110
+ )
111
+ return
vllm/models/glm5next/__init__.py
new
+16
@@ -0,0 +1,16 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from vllm.platforms import current_platform
5
+
6
+if current_platform.is_xpu():
7
+ raise NotImplementedError("GLM-5.3-Flash does not currently support XPU.")
8
+
9
+from .nvidia.model import Glm5NextForCausalLM, Glm5NextForConditionalGeneration
10
+from .nvidia.mtp import Glm5NextMTP
11
+
12
+__all__ = [
13
+ "Glm5NextForCausalLM",
14
+ "Glm5NextForConditionalGeneration",
15
+ "Glm5NextMTP",
16
+]
vllm/models/glm5next/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/glm5next/amd/ops/__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/glm5next/amd/ops/kpool_compress.py
new
+890
@@ -0,0 +1,890 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""kpool (key-pooling) Triton kernels for the sparse-attention indexer.
4
+
5
+The cache stores POOLS (1 entry per ``pool_size`` consecutive tokens) rather
6
+than individual tokens. ``compress_ratio == pool_size`` on the kv_cache_spec
7
+makes the metadata builder emit pool-granular slot_mapping / seq_lens /
8
+cu_seq_lens / page_table for free; this file supplies the compress-write
9
+kernel (replacing ``indexer_k_quant_and_cache``) and the pool-level topk
10
+helpers (select pools -> expand to tokens -> append tail).
11
+"""
12
+
13
+from __future__ import annotations
14
+
15
+import torch
16
+
17
+from vllm.platforms import current_platform
18
+from vllm.triton_utils import tl, triton
19
+
20
+# The GLM-5.3-Flash indexer head dimension is fixed at 128.
21
+INDEX_HEAD_DIM = 128
22
+FP8_DTYPE = current_platform.fp8_dtype()
23
+FP8_MAX = torch.finfo(FP8_DTYPE).max
24
+
25
+
26
+@triton.jit
27
+def _cache_k_offset(
28
+ token_offset,
29
+ dim_offset,
30
+ head_dim: tl.constexpr,
31
+ preshuffle: tl.constexpr,
32
+):
33
+ if preshuffle:
34
+ return (
35
+ (token_offset // 16) * 16 * head_dim
36
+ + (dim_offset // 16) * 16 * 16
37
+ + (token_offset % 16) * 16
38
+ + dim_offset % 16
39
+ )
40
+ return token_offset * head_dim + dim_offset
41
+
42
+
43
+# Hadamard-128 rotation
44
+
45
+
46
+@triton.jit
47
+def _hadamard128_stage(x, GROUPS: tl.constexpr, STRIDE: tl.constexpr):
48
+ x3 = tl.reshape(x, (GROUPS, 2, STRIDE))
49
+ x3 = tl.trans(x3, 0, 2, 1)
50
+ a, b = tl.split(x3)
51
+ x3 = tl.join(a + b, a - b)
52
+ x3 = tl.trans(x3, 0, 2, 1)
53
+ return tl.reshape(x3, (128,))
54
+
55
+
56
+@triton.jit
57
+def _hadamard128(x):
58
+ x = _hadamard128_stage(x, 64, 1)
59
+ x = _hadamard128_stage(x, 32, 2)
60
+ x = _hadamard128_stage(x, 16, 4)
61
+ x = _hadamard128_stage(x, 8, 8)
62
+ x = _hadamard128_stage(x, 4, 16)
63
+ x = _hadamard128_stage(x, 2, 32)
64
+ x = _hadamard128_stage(x, 1, 64)
65
+ return x * 0.08838834764831845 # 1/sqrt(128)
66
+
67
+
68
+# Map pool ids to physical cache slots.
69
+
70
+
71
+def compute_pooled_write_locs(
72
+ page_table_64: torch.Tensor,
73
+ pool_ids: torch.Tensor,
74
+ pool_size: int,
75
+) -> torch.Tensor:
76
+ """Map logical pooled-K ids to physical flat cache slots.
77
+
78
+ ``pool_size`` consecutive tokens share one pool slot that lives at the
79
+ *first* token page of each page-group. ``page_table_64`` maps token pages
80
+ to physical block ids; we gather the block id of each pool's page-group
81
+ and add the in-block pool offset.
82
+ """
83
+ assert page_table_64.ndim == 1
84
+ pool_ids = pool_ids.to(torch.int64)
85
+ block_size = 64
86
+ pool_page_group = torch.div(pool_ids, block_size, rounding_mode="floor")
87
+ token_page_row = pool_page_group * pool_size
88
+ packed_page = page_table_64.index_select(0, token_page_row.to(torch.int64))
89
+ return packed_page.to(torch.int64) * block_size + torch.remainder(
90
+ pool_ids, block_size
91
+ )
92
+
93
+
94
+def build_pooled_page_table(
95
+ page_table: torch.Tensor,
96
+ pool_size: int,
97
+) -> torch.Tensor:
98
+ """Build a pool-granular page table by taking every ``pool_size``-th
99
+ token-page column (one pool maps to ``pool_size`` token pages).
100
+
101
+ Uses gather (not strided slicing) so the result is always a fresh
102
+ row-major tensor — some downstream kernels require stride(-1) == 1.
103
+ """
104
+ block_size = page_table.shape[-1]
105
+ assert block_size % pool_size == 0, (
106
+ f"pool_size ({pool_size}) must divide page columns ({block_size})"
107
+ )
108
+ idx = torch.arange(0, block_size, pool_size, device=page_table.device)
109
+ return page_table[..., idx].contiguous()
110
+
111
+
112
+# Fused pool compression and cache write.
113
+
114
+
115
+@triton.jit
116
+def _kpool_softmax_rotate_write_cache_kernel(
117
+ buf_fp8_ptr,
118
+ buf_fp32_ptr,
119
+ slot_k_ptr,
120
+ slot_score_ptr,
121
+ ape_ptr,
122
+ loc_ptr,
123
+ write_mask_ptr,
124
+ compressed_k_ptr,
125
+ compressed_scale_ptr,
126
+ slot_k_stride_0,
127
+ slot_k_stride_1,
128
+ slot_score_stride_0,
129
+ slot_score_stride_1,
130
+ ape_stride_0,
131
+ PAGE_SIZE: tl.constexpr,
132
+ BUF_NUMEL_PER_PAGE: tl.constexpr,
133
+ POOL_SIZE: tl.constexpr,
134
+ HEAD_DIM: tl.constexpr,
135
+ S_OFFSET_NBYTES_IN_PAGE: tl.constexpr,
136
+ FP8_MAX: tl.constexpr,
137
+ PRESHUFFLE: tl.constexpr,
138
+ ROUND_SCALE: tl.constexpr,
139
+ HAS_WRITE_MASK: tl.constexpr,
140
+ RETURN_COMPRESSED: tl.constexpr,
141
+ WRITE_CACHE: tl.constexpr,
142
+ BLOCK_D: tl.constexpr,
143
+):
144
+ """One program per pool. softmax(slot_score+ape)-weighted sum of slot_k ->
145
+ Hadamard-128 -> per-vector fp8 absmax quant -> write to cache at ``loc``."""
146
+ row = tl.program_id(0)
147
+ do_write = True
148
+ if HAS_WRITE_MASK:
149
+ do_write = tl.load(write_mask_ptr + row)
150
+
151
+ offs = tl.arange(0, BLOCK_D)
152
+ mask = (offs < HEAD_DIM) & do_write
153
+
154
+ # --- Pass 1: per-dim max over the pool (softmax numerical stability) ---
155
+ max_score = tl.full((BLOCK_D,), -float("inf"), tl.float32)
156
+ for slot in tl.static_range(0, POOL_SIZE):
157
+ score = tl.load(
158
+ slot_score_ptr
159
+ + row * slot_score_stride_0
160
+ + slot * slot_score_stride_1
161
+ + offs,
162
+ mask=mask,
163
+ other=0.0,
164
+ ).to(tl.float32)
165
+ score += tl.load(ape_ptr + slot * ape_stride_0 + offs, mask=mask, other=0.0).to(
166
+ tl.float32
167
+ )
168
+ max_score = tl.maximum(max_score, score)
169
+
170
+ # --- Pass 2: softmax-weighted sum of K ---
171
+ acc = tl.full((BLOCK_D,), 0.0, tl.float32)
172
+ denom = tl.full((BLOCK_D,), 0.0, tl.float32)
173
+ for slot in tl.static_range(0, POOL_SIZE):
174
+ score = tl.load(
175
+ slot_score_ptr
176
+ + row * slot_score_stride_0
177
+ + slot * slot_score_stride_1
178
+ + offs,
179
+ mask=mask,
180
+ other=0.0,
181
+ ).to(tl.float32)
182
+ score += tl.load(ape_ptr + slot * ape_stride_0 + offs, mask=mask, other=0.0).to(
183
+ tl.float32
184
+ )
185
+ prob = tl.exp(score - max_score)
186
+ denom += prob
187
+ k = tl.load(
188
+ slot_k_ptr + row * slot_k_stride_0 + slot * slot_k_stride_1 + offs,
189
+ mask=mask,
190
+ other=0.0,
191
+ ).to(tl.float32)
192
+ acc += k * prob
193
+
194
+ x = acc / denom
195
+ x = tl.where(do_write, x, 0.0).to(tl.bfloat16).to(tl.float32)
196
+
197
+ # Match the unfused pooled-K path's bf16 precision before quantization.
198
+ x = _hadamard128(x).to(tl.bfloat16).to(tl.float32)
199
+
200
+ # --- per-vector absmax fp8 quant ---
201
+ fp8_max_inv = 1.0 / FP8_MAX
202
+ absmax = tl.max(tl.abs(x), axis=0)
203
+ absmax = tl.maximum(absmax, 1e-4)
204
+ if ROUND_SCALE:
205
+ scale = tl.exp2(tl.ceil(tl.log2(absmax * fp8_max_inv)))
206
+ else:
207
+ scale = absmax * fp8_max_inv
208
+ quantized = x / scale
209
+ quantized = tl.minimum(tl.maximum(quantized, -FP8_MAX), FP8_MAX)
210
+
211
+ if WRITE_CACHE:
212
+ loc = tl.load(loc_ptr + row, mask=do_write, other=0)
213
+ loc_page_index = loc // PAGE_SIZE
214
+ loc_token_offset_in_page = loc % PAGE_SIZE
215
+ out_k_offsets = loc_page_index * BUF_NUMEL_PER_PAGE + _cache_k_offset(
216
+ loc_token_offset_in_page,
217
+ offs,
218
+ HEAD_DIM,
219
+ PRESHUFFLE,
220
+ )
221
+ out_s_offset = (
222
+ loc_page_index * BUF_NUMEL_PER_PAGE // 4
223
+ + S_OFFSET_NBYTES_IN_PAGE // 4
224
+ + loc_token_offset_in_page
225
+ )
226
+ tl.store(buf_fp8_ptr + out_k_offsets, quantized, mask=mask)
227
+ tl.store(buf_fp32_ptr + out_s_offset, scale, mask=do_write)
228
+
229
+ if RETURN_COMPRESSED:
230
+ tl.store(
231
+ compressed_k_ptr + row * HEAD_DIM + offs,
232
+ quantized,
233
+ mask=offs < HEAD_DIM,
234
+ )
235
+ tl.store(compressed_scale_ptr + row, scale)
236
+
237
+
238
+def kpool_compress_and_write_cache(
239
+ kv_cache: torch.Tensor,
240
+ slot_k: torch.Tensor,
241
+ slot_score: torch.Tensor,
242
+ ape: torch.Tensor,
243
+ loc: torch.Tensor,
244
+ pool_size: int,
245
+ head_dim: int = INDEX_HEAD_DIM,
246
+ write_mask: torch.Tensor | None = None,
247
+ round_scale: bool = True,
248
+ return_compressed: bool = False,
249
+ write_cache: bool = True,
250
+):
251
+ """Compress ``pool_size`` tokens into one fp8 K and write at ``loc``.
252
+
253
+ Args:
254
+ kv_cache: indexer K cache ``[num_blocks, block_size, head_dim+4]`` uint8.
255
+ slot_k: ``[n_pools, pool_size, head_dim]`` bf16 — raw per-token K.
256
+ slot_score: ``[n_pools, pool_size, head_dim]`` — per-token gate score.
257
+ ape: ``[pool_size, head_dim]`` fp32 — per-slot position bias.
258
+ loc: ``[n_pools]`` int64 — flat physical slot per pool.
259
+ """
260
+ assert slot_k.ndim == 3
261
+ assert slot_score.shape == slot_k.shape
262
+ assert ape.shape == slot_k.shape[1:]
263
+ assert slot_k.shape[2] == head_dim
264
+ assert slot_k.dtype == torch.bfloat16
265
+ assert ape.dtype == torch.float32
266
+ assert kv_cache.dtype == torch.uint8
267
+ assert loc.dtype == torch.int64
268
+ assert write_cache or return_compressed
269
+
270
+ page_size = kv_cache.shape[1]
271
+ buf = kv_cache
272
+ slot_k = slot_k.contiguous()
273
+ slot_score = slot_score.contiguous()
274
+ ape = ape.contiguous()
275
+ loc = loc.contiguous()
276
+ if write_mask is None:
277
+ write_mask = torch.empty((1,), dtype=torch.bool, device=slot_k.device)
278
+ has_write_mask = False
279
+ else:
280
+ assert write_mask.shape == (slot_k.shape[0],)
281
+ write_mask = write_mask.contiguous()
282
+ has_write_mask = True
283
+ assert not return_compressed
284
+
285
+ if slot_k.shape[0] == 0:
286
+ if return_compressed:
287
+ return (
288
+ torch.empty(
289
+ (0, head_dim),
290
+ dtype=FP8_DTYPE,
291
+ device=slot_k.device,
292
+ ),
293
+ torch.empty((0,), dtype=torch.float32, device=slot_k.device),
294
+ )
295
+ return None
296
+
297
+ buf_fp8 = buf.view(FP8_DTYPE)
298
+ buf_fp32 = buf.view(torch.float32)
299
+ # bytes per page (last dim of kv_cache) viewed as uint8
300
+ buf_numel_per_page = buf.stride(0)
301
+ s_offset_nbytes_in_page = page_size * head_dim
302
+
303
+ if return_compressed:
304
+ compressed_k = torch.empty(
305
+ (slot_k.shape[0], head_dim),
306
+ dtype=FP8_DTYPE,
307
+ device=slot_k.device,
308
+ )
309
+ compressed_scale = torch.empty(
310
+ (slot_k.shape[0],), dtype=torch.float32, device=slot_k.device
311
+ )
312
+ else:
313
+ compressed_k = buf_fp8
314
+ compressed_scale = buf_fp32
315
+
316
+ if page_size > 1:
317
+ assert page_size % 16 == 0, "ROCm preshuffle requires 16-token tiles"
318
+
319
+ _kpool_softmax_rotate_write_cache_kernel[(slot_k.shape[0],)](
320
+ buf_fp8,
321
+ buf_fp32,
322
+ slot_k,
323
+ slot_score,
324
+ ape,
325
+ loc,
326
+ write_mask,
327
+ compressed_k,
328
+ compressed_scale,
329
+ slot_k.stride(0),
330
+ slot_k.stride(1),
331
+ slot_score.stride(0),
332
+ slot_score.stride(1),
333
+ ape.stride(0),
334
+ PAGE_SIZE=page_size,
335
+ BUF_NUMEL_PER_PAGE=buf_numel_per_page,
336
+ POOL_SIZE=slot_k.shape[1],
337
+ HEAD_DIM=head_dim,
338
+ S_OFFSET_NBYTES_IN_PAGE=s_offset_nbytes_in_page,
339
+ FP8_MAX=FP8_MAX,
340
+ PRESHUFFLE=page_size > 1,
341
+ ROUND_SCALE=round_scale,
342
+ HAS_WRITE_MASK=has_write_mask,
343
+ RETURN_COMPRESSED=return_compressed,
344
+ WRITE_CACHE=write_cache,
345
+ BLOCK_D=triton.next_power_of_2(head_dim),
346
+ )
347
+
348
+ if return_compressed:
349
+ return compressed_k, compressed_scale
350
+ return None
351
+
352
+
353
+# Seed each request's incomplete pool into its paged tail during prefill.
354
+
355
+
356
+@triton.jit
357
+def _kpool_tail_seed_kernel(
358
+ key_ptr,
359
+ score_ptr,
360
+ tslot_ptr,
361
+ tail_ptr,
362
+ n_tokens,
363
+ TAIL_BLOCK_ELEMS: tl.constexpr,
364
+ KPOOL_HEAD: tl.constexpr,
365
+ HEAD_DIM: tl.constexpr,
366
+ KPOOL: tl.constexpr,
367
+ BLOCK_D: tl.constexpr,
368
+):
369
+ """Copy token ``i``'s raw K + gate into its request's tail block.
370
+
371
+ Token ``i`` is among its request's last KPOOL tokens iff the token KPOOL
372
+ ahead belongs to a different tail block (or is past the batch / padding,
373
+ slot < 0). ``tslot = block * KPOOL + pos % KPOOL``; the destination is
374
+ ``tail[block, {0:K, 1:score}, pos % KPOOL, :]``.
375
+ """
376
+ i = tl.program_id(0)
377
+ t = tl.load(tslot_ptr + i).to(tl.int64)
378
+ if t < 0:
379
+ return
380
+ blk = t // KPOOL # t >= 0 here, so trunc == floor
381
+ ahead = tl.load(tslot_ptr + i + KPOOL, mask=i + KPOOL < n_tokens, other=-1).to(
382
+ tl.int64
383
+ )
384
+ # Match the torch semantics exactly: a negative ahead slot floors to a
385
+ # block id that differs from every real block -> token is in the tail.
386
+ # Only divide non-negative slots (Triton int div truncates, torch floors).
387
+ if ahead >= 0 and ahead // KPOOL == blk:
388
+ return
389
+ offs = tl.arange(0, BLOCK_D)
390
+ m = offs < HEAD_DIM
391
+ block_base = blk * TAIL_BLOCK_ELEMS
392
+ base = block_base + (t % KPOOL) * HEAD_DIM
393
+ k = tl.load(key_ptr + i * HEAD_DIM + offs, mask=m)
394
+ s = tl.load(score_ptr + i * HEAD_DIM + offs, mask=m)
395
+ tl.store(tail_ptr + base + offs, k, mask=m)
396
+ tl.store(
397
+ tail_ptr + block_base + KPOOL_HEAD + (t % KPOOL) * HEAD_DIM + offs, s, mask=m
398
+ )
399
+
400
+
401
+def kpool_seed_tail_cache(
402
+ tail_kv_cache: torch.Tensor,
403
+ key: torch.Tensor,
404
+ gate_score: torch.Tensor,
405
+ tslot: torch.Tensor,
406
+ kpool: int,
407
+ head_dim: int = INDEX_HEAD_DIM,
408
+) -> None:
409
+ """Seed the paged tail cache from a prefill batch (see the kernel)."""
410
+ assert tail_kv_cache.dtype == torch.bfloat16
411
+ assert key.dtype == torch.bfloat16
412
+ n = tslot.shape[0]
413
+ if n == 0:
414
+ return
415
+ _kpool_tail_seed_kernel[(n,)](
416
+ key,
417
+ gate_score,
418
+ tslot,
419
+ tail_kv_cache,
420
+ n,
421
+ TAIL_BLOCK_ELEMS=tail_kv_cache.stride(0),
422
+ KPOOL_HEAD=tail_kv_cache.stride(1),
423
+ HEAD_DIM=head_dim,
424
+ KPOOL=kpool,
425
+ BLOCK_D=triton.next_power_of_2(head_dim),
426
+ )
427
+
428
+
429
+# Update each request's tail during decode and write completed pools.
430
+
431
+
432
+@triton.jit
433
+def _kpool_decode_update_batched_kernel(
434
+ buf_fp8_ptr,
435
+ buf_fp32_ptr,
436
+ tail_kv_ptr,
437
+ tail_slot_mapping_ptr, # [B, NEXT_N] int32
438
+ key_ptr, # [B, NEXT_N, HEAD_DIM] bf16
439
+ key_stride_b,
440
+ key_stride_t,
441
+ slot_score_ptr, # [B, NEXT_N, HEAD_DIM] bf16
442
+ ss_stride_b,
443
+ ss_stride_t,
444
+ ape_ptr,
445
+ ape_stride_0,
446
+ slot_mapping_ptr, # [B, NEXT_N] int32
447
+ positions_ptr, # [B, NEXT_N] int32
448
+ NEXT_N, # runtime token count per request (no .item() needed)
449
+ PAGE_SIZE: tl.constexpr,
450
+ BUF_NUMEL_PER_PAGE: tl.constexpr,
451
+ POOL_SIZE: tl.constexpr,
452
+ TAIL_BLOCK_ELEMS: tl.constexpr,
453
+ KPOOL_HEAD: tl.constexpr,
454
+ HEAD_DIM: tl.constexpr,
455
+ S_OFFSET_NBYTES_IN_PAGE: tl.constexpr,
456
+ FP8_MAX: tl.constexpr,
457
+ PRESHUFFLE: tl.constexpr,
458
+ ROUND_SCALE: tl.constexpr,
459
+ BLOCK_D: tl.constexpr,
460
+):
461
+ """One program per request; iterates its NEXT_N verify tokens in order.
462
+
463
+ Replaces the caller's per-token sequential launch loop. The intra-request
464
+ iteration MUST stay in position order: a pool-completion at token t* reads
465
+ the tail-ring slots that tokens t < t* (same request) just stashed in this
466
+ same invocation. ``tl.range`` iterates sequentially within the program, so
467
+ those stashes are visible to the later completion read. Cross-request
468
+ programs are independent (distinct tail blocks). With NEXT_N < POOL_SIZE
469
+ (the spec-verify case: NEXT_N ~= num_spec+1, POOL_SIZE=16) at most one
470
+ completion can occur per request per call, but the ordered loop is correct
471
+ for any NEXT_N.
472
+ """
473
+ req = tl.program_id(0)
474
+ offs = tl.arange(0, BLOCK_D)
475
+ dim_mask = offs < HEAD_DIM
476
+
477
+ for t in tl.range(0, NEXT_N):
478
+ idx = req * NEXT_N + t
479
+ cache_loc = tl.load(slot_mapping_ptr + idx)
480
+ pos = tl.load(positions_ptr + idx)
481
+ safe_pos = tl.maximum(pos, 0)
482
+ pos_valid = (cache_loc >= 0) & (pos >= 0)
483
+
484
+ slot = safe_pos % POOL_SIZE
485
+ phys_slot = safe_pos % POOL_SIZE
486
+
487
+ # Derive the tail block from THIS token's tail_slot (the request's block
488
+ # is constant across a pool, but a padded / invalid entry carries a
489
+ # negative sentinel -- reading it from token 0 would poison every
490
+ # token's base address). Clamp so an invalid entry can never form an
491
+ # out-of-bounds base; the accesses below are gated on pos_valid anyway.
492
+ tail_slot = tl.load(tail_slot_mapping_ptr + idx)
493
+ block = tl.maximum(tail_slot, 0).to(tl.int64) // POOL_SIZE
494
+ block_base = block * TAIL_BLOCK_ELEMS
495
+
496
+ # The tail-ring stash must run for EVERY real token, so it is gated on
497
+ # the token-granular tail slot -- not on `pos_valid`, which keys off the
498
+ # POOL-granular `slot_mapping` and is therefore only true on the pool's
499
+ # last token. Gating the stash on pos_valid dropped every intra-pool
500
+ # token, so a decode-built pool compressed 3 stale ring entries (the
501
+ # prefill-seeded prompt tail, frozen forever) plus the current token.
502
+ stash_valid = (pos >= 0) & (tail_slot >= 0)
503
+
504
+ key = tl.load(
505
+ key_ptr + req * key_stride_b + t * key_stride_t + offs,
506
+ mask=dim_mask,
507
+ other=0.0,
508
+ ).to(tl.float32)
509
+ score_current = tl.load(
510
+ slot_score_ptr + req * ss_stride_b + t * ss_stride_t + offs,
511
+ mask=dim_mask,
512
+ other=0.0,
513
+ ).to(tl.float32)
514
+
515
+ if pos_valid & (slot == POOL_SIZE - 1):
516
+ pool_logical_start = safe_pos - slot
517
+
518
+ max_score = tl.full((BLOCK_D,), -float("inf"), tl.float32)
519
+ for pool_slot in tl.static_range(0, POOL_SIZE):
520
+ is_current = pool_slot == slot
521
+ phys = (pool_logical_start + pool_slot) % POOL_SIZE
522
+ score_buf = tl.load(
523
+ tail_kv_ptr + block_base + KPOOL_HEAD + phys * HEAD_DIM + offs,
524
+ mask=dim_mask,
525
+ other=0.0,
526
+ ).to(tl.float32)
527
+ score = tl.where(is_current, score_current, score_buf)
528
+ score += tl.load(
529
+ ape_ptr + pool_slot * ape_stride_0 + offs,
530
+ mask=dim_mask,
531
+ other=0.0,
532
+ ).to(tl.float32)
533
+ max_score = tl.maximum(max_score, score)
534
+
535
+ acc = tl.full((BLOCK_D,), 0.0, tl.float32)
536
+ denom = tl.full((BLOCK_D,), 0.0, tl.float32)
537
+ for pool_slot in tl.static_range(0, POOL_SIZE):
538
+ is_current = pool_slot == slot
539
+ phys = (pool_logical_start + pool_slot) % POOL_SIZE
540
+ score_buf = tl.load(
541
+ tail_kv_ptr + block_base + KPOOL_HEAD + phys * HEAD_DIM + offs,
542
+ mask=dim_mask,
543
+ other=0.0,
544
+ ).to(tl.float32)
545
+ score = tl.where(is_current, score_current, score_buf)
546
+ score += tl.load(
547
+ ape_ptr + pool_slot * ape_stride_0 + offs,
548
+ mask=dim_mask,
549
+ other=0.0,
550
+ ).to(tl.float32)
551
+ prob = tl.exp(score - max_score)
552
+ denom += prob
553
+ k_buf = tl.load(
554
+ tail_kv_ptr + block_base + phys * HEAD_DIM + offs,
555
+ mask=dim_mask,
556
+ other=0.0,
557
+ ).to(tl.float32)
558
+ k = tl.where(is_current, key, k_buf)
559
+ acc += k * prob
560
+
561
+ x = (acc / denom).to(tl.bfloat16).to(tl.float32)
562
+ x = _hadamard128(x).to(tl.bfloat16).to(tl.float32)
563
+
564
+ fp8_max_inv = 1.0 / FP8_MAX
565
+ absmax = tl.maximum(tl.max(tl.abs(x), axis=0), 1e-4)
566
+ if ROUND_SCALE:
567
+ scale = tl.exp2(tl.ceil(tl.log2(absmax * fp8_max_inv)))
568
+ else:
569
+ scale = absmax * fp8_max_inv
570
+ quantized = tl.minimum(tl.maximum(x / scale, -FP8_MAX), FP8_MAX)
571
+
572
+ loc = cache_loc.to(tl.int64)
573
+ loc_page_index = loc // PAGE_SIZE
574
+ loc_token_offset_in_page = loc % PAGE_SIZE
575
+ out_k_offsets = loc_page_index * BUF_NUMEL_PER_PAGE + _cache_k_offset(
576
+ loc_token_offset_in_page,
577
+ offs,
578
+ HEAD_DIM,
579
+ PRESHUFFLE,
580
+ )
581
+ out_s_offset = (
582
+ loc_page_index * BUF_NUMEL_PER_PAGE // 4
583
+ + S_OFFSET_NBYTES_IN_PAGE // 4
584
+ + loc_token_offset_in_page
585
+ )
586
+ tl.store(buf_fp8_ptr + out_k_offsets, quantized, mask=dim_mask)
587
+ tl.store(buf_fp32_ptr + out_s_offset, scale)
588
+
589
+ # Stash the current token AFTER any completion read so the completion
590
+ # uses prior stashes (and the current token's own key/score via
591
+ # is_current), then leaves this token for future pools. Order matches
592
+ # the per-token kernel: completion read first, stash second.
593
+ update_mask = dim_mask & stash_valid
594
+ tl.store(
595
+ tail_kv_ptr + block_base + phys_slot * HEAD_DIM + offs,
596
+ key,
597
+ mask=update_mask,
598
+ )
599
+ tl.store(
600
+ tail_kv_ptr + block_base + KPOOL_HEAD + phys_slot * HEAD_DIM + offs,
601
+ score_current,
602
+ mask=update_mask,
603
+ )
604
+
605
+
606
+def kpool_decode_update_and_maybe_write_cache_batched(
607
+ kv_cache: torch.Tensor,
608
+ tail_kv_cache: torch.Tensor,
609
+ tail_slot_mapping: torch.Tensor,
610
+ key: torch.Tensor,
611
+ slot_score: torch.Tensor,
612
+ ape: torch.Tensor,
613
+ slot_mapping: torch.Tensor,
614
+ positions: torch.Tensor,
615
+ pool_size: int,
616
+ head_dim: int = INDEX_HEAD_DIM,
617
+ round_scale: bool = True,
618
+) -> None:
619
+ """Batched decode-step kpool update for spec verify (``next_n > 1``).
620
+
621
+ One launch replaces the caller's per-token loop. Inputs are grouped per
622
+ request: ``[num_requests, next_n, ...]``. Each program handles one
623
+ request's ``next_n`` tokens in position order (see the kernel docstring for
624
+ why ordering is required for pool-completion correctness).
625
+
626
+ Plain decode (``next_n == 1``) is handled here too — the kernel collapses
627
+ to a single-iteration loop.
628
+
629
+ Args:
630
+ kv_cache: indexer K cache ``[num_blocks, block_size, head_dim+4]`` uint8.
631
+ tail_kv_cache: paged tail cache ``[num_blocks, 2, pool_size, head_dim]``
632
+ bf16 (K at half 0, gate score at half 1).
633
+ tail_slot_mapping: ``[num_requests, next_n]`` int32.
634
+ key / slot_score: ``[num_requests, next_n, head_dim]`` bf16.
635
+ ape: ``[pool_size, head_dim]`` fp32.
636
+ slot_mapping / positions: ``[num_requests, next_n]`` int32.
637
+ """
638
+ num_requests, next_n = key.shape[0], key.shape[1]
639
+ if num_requests == 0 or next_n == 0:
640
+ return
641
+ assert tail_kv_cache.ndim == 4
642
+ assert tail_kv_cache.shape[1] == 2
643
+ assert tail_kv_cache.shape[2] == pool_size
644
+ assert tail_kv_cache.shape[3] == head_dim
645
+ assert tail_kv_cache.dtype == torch.bfloat16
646
+ assert key.ndim == 3 and key.shape[2] == head_dim
647
+ assert slot_score.shape == key.shape
648
+ assert ape.shape == (pool_size, head_dim)
649
+ assert tail_slot_mapping.shape == (num_requests, next_n)
650
+ assert slot_mapping.shape == (num_requests, next_n)
651
+ assert positions.shape == (num_requests, next_n)
652
+ assert key.dtype == torch.bfloat16
653
+ assert slot_score.dtype == torch.bfloat16
654
+ assert ape.dtype == torch.float32
655
+ assert kv_cache.dtype == torch.uint8
656
+
657
+ page_size = kv_cache.shape[1]
658
+ buf = kv_cache
659
+ buf_fp8 = buf.view(FP8_DTYPE)
660
+ buf_fp32 = buf.view(torch.float32)
661
+
662
+ # The kernel indexes the int tensors as ``req * next_n + t`` (row-major),
663
+ # so they must be contiguous. Callers pass either a view of a contiguous
664
+ # slice or a freshly scattered tensor, making these no-ops; the calls guard
665
+ # against a future caller handing over a strided view.
666
+ tail_slot_mapping = tail_slot_mapping.contiguous()
667
+ slot_mapping = slot_mapping.contiguous()
668
+ positions = positions.contiguous()
669
+
670
+ if page_size > 1:
671
+ assert page_size % 16 == 0, "ROCm preshuffle requires 16-token tiles"
672
+
673
+ _kpool_decode_update_batched_kernel[(num_requests,)](
674
+ buf_fp8,
675
+ buf_fp32,
676
+ tail_kv_cache,
677
+ tail_slot_mapping,
678
+ key,
679
+ key.stride(0),
680
+ key.stride(1),
681
+ slot_score,
682
+ slot_score.stride(0),
683
+ slot_score.stride(1),
684
+ ape,
685
+ ape.stride(0),
686
+ slot_mapping,
687
+ positions,
688
+ next_n,
689
+ PAGE_SIZE=page_size,
690
+ BUF_NUMEL_PER_PAGE=buf.stride(0),
691
+ POOL_SIZE=pool_size,
692
+ TAIL_BLOCK_ELEMS=tail_kv_cache.stride(0),
693
+ KPOOL_HEAD=tail_kv_cache.stride(1),
694
+ HEAD_DIM=head_dim,
695
+ S_OFFSET_NBYTES_IN_PAGE=page_size * head_dim,
696
+ FP8_MAX=FP8_MAX,
697
+ PRESHUFFLE=page_size > 1,
698
+ ROUND_SCALE=round_scale,
699
+ BLOCK_D=triton.next_power_of_2(head_dim),
700
+ )
701
+
702
+
703
+# Pool-level top-k helpers.
704
+
705
+
706
+def history_group_budget_for_topk(topk: int, pool_size: int) -> int:
707
+ """Number of pools to select so that expanding yields ``topk`` tokens."""
708
+ assert topk % pool_size == 0
709
+ return topk // pool_size
710
+
711
+
712
+def expand_pools_to_tokens(
713
+ group_ids: torch.Tensor,
714
+ group_valid: torch.Tensor,
715
+ topk: int,
716
+ pool_size: int,
717
+ page_table: torch.Tensor | None = None,
718
+ topk_offsets: torch.Tensor | None = None,
719
+) -> torch.Tensor:
720
+ """Expand selected full-pool ids to a strict-width token topk tensor."""
721
+ assert group_ids.ndim == 2
722
+ assert group_valid.shape == group_ids.shape
723
+ assert topk % pool_size == 0
724
+ assert group_ids.shape[1] == history_group_budget_for_topk(topk, pool_size)
725
+ assert page_table is None or topk_offsets is None
726
+
727
+ device = group_ids.device
728
+ offsets = torch.arange(pool_size, device=device, dtype=torch.int64)
729
+ token_ids = group_ids.to(torch.int64).unsqueeze(-1) * pool_size + offsets
730
+ token_ids = token_ids.reshape(group_ids.shape[0], topk)
731
+ valid = (
732
+ group_valid.unsqueeze(-1)
733
+ .expand(-1, -1, pool_size)
734
+ .reshape(group_ids.shape[0], topk)
735
+ )
736
+
737
+ if page_table is not None:
738
+ assert page_table.ndim == 2
739
+ safe_ids = token_ids.clamp(min=0, max=page_table.shape[1] - 1)
740
+ output = torch.gather(page_table, dim=1, index=safe_ids).to(torch.int32)
741
+ elif topk_offsets is not None:
742
+ if topk_offsets.ndim == 2:
743
+ assert topk_offsets.shape[1] == 1
744
+ topk_offsets = topk_offsets.squeeze(1)
745
+ output = (token_ids + topk_offsets.to(torch.int64).unsqueeze(1)).to(torch.int32)
746
+ else:
747
+ output = token_ids.to(torch.int32)
748
+
749
+ return torch.where(valid, output, torch.full_like(output, -1))
750
+
751
+
752
+def append_tail_to_topk(
753
+ topk_result: torch.Tensor,
754
+ seq_lens: torch.Tensor,
755
+ pool_lens: torch.Tensor,
756
+ pool_size: int,
757
+ page_table: torch.Tensor | None = None,
758
+ topk_offsets: torch.Tensor | None = None,
759
+) -> torch.Tensor:
760
+ """Append non-pooled tail tokens after expanded history tokens.
761
+
762
+ ``index_kpool_always_select_tail`` keeps the (incomplete) trailing pool so
763
+ the most recent tokens are always attended to.
764
+ """
765
+ assert topk_result.dtype == torch.int32
766
+ assert seq_lens.ndim == 1
767
+ assert pool_lens.ndim == 1
768
+
769
+ tail_pool = pool_size - 1
770
+ if tail_pool == 0:
771
+ return topk_result
772
+
773
+ rows, n_cols = topk_result.shape
774
+ out_cols = n_cols + tail_pool
775
+ out = torch.empty(
776
+ (rows, out_cols), dtype=topk_result.dtype, device=topk_result.device
777
+ )
778
+
779
+ # tail tokens: [pool_len*pool_size, seq_len) for each row.
780
+ pool_len = pool_lens.to(torch.int32)
781
+ tail_start = pool_len * pool_size
782
+ seq_len = seq_lens.to(torch.int32)
783
+ tail_count = seq_len - tail_start # in [0, pool_size)
784
+
785
+ cols = torch.arange(out_cols, device=topk_result.device)[None, :]
786
+ history_len = n_cols
787
+ is_history = cols < history_len
788
+ tail_off = cols - history_len
789
+ is_tail = (tail_off >= 0) & (tail_off < tail_count[:, None])
790
+
791
+ # safe_hist must be per-row [rows, out_cols] so the gather reads each row's
792
+ # OWN history. cols is [1, out_cols]; if used directly, gather (which does
793
+ # NOT broadcast the index) would read only row 0 of topk_result, making every
794
+ # query inherit row 0's history (empty for the first token) and lose all its
795
+ # selected tokens — only the per-row tail would survive. This only manifests
796
+ # for multi-row sparse PREFILL (decode has 1 row, so it reads its own row 0).
797
+ safe_hist = torch.minimum(cols, torch.full_like(cols, n_cols - 1)).expand(
798
+ rows, out_cols
799
+ )
800
+ history_val = torch.gather(topk_result, 1, safe_hist)
801
+
802
+ tail_raw = tail_start[:, None] + tail_off
803
+ tail_val = tail_raw.to(torch.int32)
804
+ if page_table is not None:
805
+ safe_tail = tail_raw.clamp(min=0, max=page_table.shape[1] - 1)
806
+ tail_val = torch.gather(page_table, 1, safe_tail).to(torch.int32)
807
+ elif topk_offsets is not None:
808
+ tail_val = (tail_raw + topk_offsets.to(torch.int64).unsqueeze(1)).to(
809
+ torch.int32
810
+ )
811
+
812
+ out = torch.where(is_history, history_val, -1)
813
+ out = torch.where(is_tail, tail_val, out)
814
+ return out
815
+
816
+
817
+@triton.jit
818
+def _expand_pools_and_append_tail_kernel(
819
+ pool_ids_ptr, # [rows, n_groups], int (any int dtype)
820
+ seq_lens_ptr, # [rows], int32 (token-granular seq_len)
821
+ out_ptr, # [rows, out_cols], int32
822
+ topk, # n_groups * pool_size
823
+ out_cols, # topk + pool_size - 1
824
+ POOL_SIZE: tl.constexpr,
825
+ BLOCK_COLS: tl.constexpr,
826
+ pid_s0,
827
+ out_s0,
828
+):
829
+ # Fuses expand_pools_to_tokens + append_tail_to_topk (identity path) into a
830
+ # single kernel. Each program writes one (row, column-tile) of the output.
831
+ row = tl.program_id(0)
832
+ tile = tl.program_id(1)
833
+ cols = tile * BLOCK_COLS + tl.arange(0, BLOCK_COLS)
834
+ mask = cols < out_cols
835
+
836
+ seq_len = tl.load(seq_lens_ptr + row)
837
+ pool_len = seq_len // POOL_SIZE
838
+ tail_start = pool_len * POOL_SIZE
839
+ tail_count = seq_len - tail_start # in [0, POOL_SIZE)
840
+
841
+ # History region [0, topk): expand selected pool g = cols // POOL_SIZE.
842
+ is_history = cols < topk
843
+ g = cols // POOL_SIZE
844
+ o = cols % POOL_SIZE
845
+ pid = tl.load(pool_ids_ptr + row * pid_s0 + g, mask=mask & is_history, other=-1)
846
+ hist_val = (pid * POOL_SIZE + o).to(tl.int32)
847
+ hist_out = tl.where(pid >= 0, hist_val, -1)
848
+
849
+ # Tail region [topk, out_cols): the request's trailing incomplete pool.
850
+ tail_off = cols - topk
851
+ is_tail = (tail_off >= 0) & (tail_off < tail_count)
852
+ tail_val = (tail_start + tail_off).to(tl.int32)
853
+ tail_out = tl.where(is_tail, tail_val, -1)
854
+
855
+ result = tl.where(is_history, hist_out, tail_out)
856
+ tl.store(out_ptr + row * out_s0 + cols, result, mask=mask)
857
+
858
+
859
+def expand_pools_and_append_tail(
860
+ pool_ids: torch.Tensor,
861
+ seq_lens: torch.Tensor,
862
+ pool_size: int,
863
+) -> torch.Tensor:
864
+ """Fuse ``expand_pools_to_tokens`` + ``append_tail_to_topk`` (identity path).
865
+
866
+ Produces the same ``[rows, topk + pool_size - 1]`` int32 output as calling
867
+ the two functions in sequence when neither ``page_table`` nor
868
+ ``topk_offsets`` is passed — the only path used by the GLM-5.3-Flash indexer.
869
+ The kernel derives ``pool_len = seq_len // pool_size`` internally, so the
870
+ caller no longer needs to precompute it. Replaces ~25 elementwise kernels
871
+ with one Triton launch.
872
+ """
873
+ rows, n_groups = pool_ids.shape
874
+ topk = n_groups * pool_size
875
+ out_cols = topk + pool_size - 1
876
+ out = torch.empty((rows, out_cols), dtype=torch.int32, device=pool_ids.device)
877
+ BLOCK_COLS = 128
878
+ n_tiles = triton.cdiv(out_cols, BLOCK_COLS)
879
+ _expand_pools_and_append_tail_kernel[(rows, n_tiles)](
880
+ pool_ids,
881
+ seq_lens,
882
+ out,
883
+ topk,
884
+ out_cols,
885
+ POOL_SIZE=pool_size,
886
+ BLOCK_COLS=BLOCK_COLS,
887
+ pid_s0=pool_ids.stride(0),
888
+ out_s0=out.stride(0),
889
+ )
890
+ return out
vllm/models/glm5next/amd/ops/third_party/__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/glm5next/amd/ops/third_party/kda/__init__.py
new
+6
@@ -0,0 +1,6 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from .kernels import chunk_kda_with_fused_gate, fused_recurrent_kda
5
+
6
+__all__ = ["chunk_kda_with_fused_gate", "fused_recurrent_kda"]
vllm/models/glm5next/amd/ops/third_party/kda/fused_recurrent.py
new
+656
@@ -0,0 +1,656 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang
4
+# mypy: ignore-errors
5
+#
6
+# This file contains code copied from the flash-linear-attention project.
7
+# The original source code was licensed under the MIT license and included
8
+# the following copyright notice:
9
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
10
+# ruff: noqa: E501
11
+
12
+import torch
13
+
14
+from vllm.third_party.flash_linear_attention.ops.op import exp
15
+from vllm.triton_utils import tl, triton
16
+
17
+
18
+@triton.heuristics(
19
+ {
20
+ "USE_INITIAL_STATE": lambda args: args["h0"] is not None,
21
+ "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
22
+ "IS_CONTINUOUS_BATCHING": lambda args: args["ssm_state_indices"] is not None,
23
+ "IS_SPEC_DECODING": lambda args: args["num_accepted_tokens"] is not None,
24
+ }
25
+)
26
+@triton.jit(do_not_specialize=["N", "T"])
27
+def fused_recurrent_gated_delta_rule_fwd_kernel(
28
+ q,
29
+ k,
30
+ v,
31
+ g,
32
+ beta,
33
+ o,
34
+ h0,
35
+ ht,
36
+ cu_seqlens,
37
+ ssm_state_indices,
38
+ num_accepted_tokens,
39
+ a_log,
40
+ g_bias,
41
+ scale,
42
+ N: tl.int64, # num of sequences
43
+ T: tl.int64, # num of tokens
44
+ B: tl.constexpr,
45
+ H: tl.constexpr,
46
+ HV: tl.constexpr,
47
+ K: tl.constexpr,
48
+ V: tl.constexpr,
49
+ BK: tl.constexpr,
50
+ BV: tl.constexpr,
51
+ stride_init_state_token: tl.constexpr,
52
+ stride_final_state_token: tl.constexpr,
53
+ stride_indices_seq: tl.constexpr,
54
+ stride_indices_tok: tl.constexpr,
55
+ USE_INITIAL_STATE: tl.constexpr, # whether to use initial state
56
+ INPLACE_FINAL_STATE: tl.constexpr, # whether to store final state inplace
57
+ IS_BETA_HEADWISE: tl.constexpr, # whether beta is headwise vector or scalar,
58
+ USE_QK_L2NORM_IN_KERNEL: tl.constexpr,
59
+ IS_VARLEN: tl.constexpr,
60
+ IS_CONTINUOUS_BATCHING: tl.constexpr,
61
+ IS_SPEC_DECODING: tl.constexpr,
62
+ IS_KDA: tl.constexpr,
63
+ SIGMOID_BETA: tl.constexpr, # beta holds raw logits; sigmoid at fp32 load
64
+ COMPUTE_GATE: tl.constexpr, # g holds raw logits; KDA gate computed in-kernel
65
+ SAFE_GATE: tl.constexpr, # bounded gate variant (only branch implemented)
66
+ LOWER_BOUND: tl.constexpr,
67
+):
68
+ i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
69
+ i_n, i_hv = i_nh // HV, i_nh % HV
70
+ i_h = i_hv // (HV // H)
71
+ if IS_VARLEN:
72
+ bos, eos = (
73
+ tl.load(cu_seqlens + i_n).to(tl.int64),
74
+ tl.load(cu_seqlens + i_n + 1).to(tl.int64),
75
+ )
76
+ all = T
77
+ T = eos - bos
78
+ else:
79
+ bos, eos = i_n * T, i_n * T + T
80
+ all = B * T
81
+
82
+ if T == 0:
83
+ # no tokens to process for this sequence
84
+ return
85
+
86
+ o_k = i_k * BK + tl.arange(0, BK)
87
+ o_v = i_v * BV + tl.arange(0, BV)
88
+
89
+ p_q = q + (bos * H + i_h) * K + o_k
90
+ p_k = k + (bos * H + i_h) * K + o_k
91
+ p_v = v + (bos * HV + i_hv) * V + o_v
92
+ if IS_BETA_HEADWISE:
93
+ p_beta = beta + (bos * HV + i_hv) * V + o_v
94
+ else:
95
+ p_beta = beta + bos * HV + i_hv
96
+
97
+ if not IS_KDA:
98
+ p_g = g + bos * HV + i_hv
99
+ else:
100
+ p_gk = g + (bos * HV + i_hv) * K + o_k
101
+
102
+ # Per-head gate amplitude, hoisted out of the token loop (COMPUTE_GATE).
103
+ if COMPUTE_GATE:
104
+ b_a_log = tl.exp(tl.load(a_log + i_h).to(tl.float32))
105
+
106
+ p_o = o + ((i_k * all + bos) * HV + i_hv) * V + o_v
107
+
108
+ mask_k = o_k < K
109
+ mask_v = o_v < V
110
+ mask_h = mask_v[:, None] & mask_k[None, :]
111
+
112
+ b_h = tl.zeros([BV, BK], dtype=tl.float32)
113
+ if USE_INITIAL_STATE:
114
+ if IS_CONTINUOUS_BATCHING:
115
+ if IS_SPEC_DECODING:
116
+ i_t = tl.load(num_accepted_tokens + i_n).to(tl.int64) - 1
117
+ else:
118
+ i_t = 0
119
+ # Load state index and check for invalid entries
120
+ state_idx = tl.load(ssm_state_indices + i_n * stride_indices_seq + i_t).to(
121
+ tl.int64
122
+ )
123
+ # Skip if state index is invalid (NULL_BLOCK_ID=0)
124
+ if state_idx <= 0:
125
+ return
126
+ p_h0 = h0 + state_idx * stride_init_state_token
127
+ else:
128
+ p_h0 = h0 + bos * HV * V * K
129
+ p_h0 = p_h0 + i_hv * V * K + o_v[:, None] * K + o_k[None, :]
130
+ b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32)
131
+
132
+ for i_t in range(0, T):
133
+ b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32)
134
+ b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32)
135
+ b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32)
136
+
137
+ if USE_QK_L2NORM_IN_KERNEL:
138
+ b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6)
139
+ b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6)
140
+ b_q = b_q * scale
141
+ # [BV, BK]
142
+ if not IS_KDA:
143
+ b_g = tl.load(p_g).to(tl.float32)
144
+ b_h *= exp(b_g)
145
+ else:
146
+ b_gk = tl.load(p_gk).to(tl.float32)
147
+ if COMPUTE_GATE:
148
+ # Replicates kda_gate_fwd_kernel's SAFE_GATE branch
149
+ # bit-for-bit (same tl.exp, same fp32 math; the intermediate
150
+ # gate value this replaces was stored/reloaded as fp32,
151
+ # which is lossless): y = lb / (1 + exp(-exp(A)*(g+bias))).
152
+ b_gk += tl.load(g_bias + i_h * K + o_k, mask=mask_k, other=0.0).to(
153
+ tl.float32
154
+ )
155
+ b_gk = LOWER_BOUND / (1.0 + tl.exp(-(b_a_log * b_gk)))
156
+ b_h *= exp(b_gk[None, :])
157
+ # [BV]
158
+ b_v -= tl.sum(b_h * b_k[None, :], 1)
159
+ if IS_BETA_HEADWISE:
160
+ b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32)
161
+ else:
162
+ b_beta = tl.load(p_beta).to(tl.float32)
163
+ # Matches torch's `x.float().sigmoid()` pre-computation bit-for-bit
164
+ # on the input side (bf16->fp32 is exact); only the sigmoid impl itself
165
+ # can differ by <=1 ULP.
166
+ if SIGMOID_BETA:
167
+ b_beta = tl.sigmoid(b_beta)
168
+ b_v *= b_beta
169
+ # [BV, BK]
170
+ b_h += b_v[:, None] * b_k[None, :]
171
+ # [BV]
172
+ b_o = tl.sum(b_h * b_q[None, :], 1)
173
+ tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v)
174
+
175
+ # keep the states for multi-query tokens
176
+ if INPLACE_FINAL_STATE:
177
+ # Load state index and check for invalid entries
178
+ final_state_idx = tl.load(
179
+ ssm_state_indices + i_n * stride_indices_seq + i_t
180
+ ).to(tl.int64)
181
+ # Only store if state index is valid (not NULL_BLOCK_ID=0)
182
+ if final_state_idx > 0:
183
+ p_ht = ht + final_state_idx * stride_final_state_token
184
+ p_ht = p_ht + i_hv * V * K + o_v[:, None] * K + o_k[None, :]
185
+ tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h)
186
+ else:
187
+ p_ht = ht + (bos + i_t) * stride_final_state_token
188
+ p_ht = p_ht + i_hv * V * K + o_v[:, None] * K + o_k[None, :]
189
+ tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h)
190
+
191
+ p_q += H * K
192
+ p_k += H * K
193
+ p_o += HV * V
194
+ p_v += HV * V
195
+ if not IS_KDA:
196
+ p_g += HV
197
+ else:
198
+ p_gk += HV * K
199
+ p_beta += HV * (V if IS_BETA_HEADWISE else 1)
200
+
201
+
202
+def fused_recurrent_gated_delta_rule_fwd(
203
+ q: torch.Tensor,
204
+ k: torch.Tensor,
205
+ v: torch.Tensor,
206
+ g: torch.Tensor,
207
+ beta: torch.Tensor,
208
+ scale: float,
209
+ initial_state: torch.Tensor,
210
+ inplace_final_state: bool = True,
211
+ cu_seqlens: torch.Tensor | None = None,
212
+ ssm_state_indices: torch.Tensor | None = None,
213
+ num_accepted_tokens: torch.Tensor | None = None,
214
+ use_qk_l2norm_in_kernel: bool = False,
215
+) -> tuple[torch.Tensor, torch.Tensor]:
216
+ B, T, H, K, V = *k.shape, v.shape[-1]
217
+ HV = v.shape[2]
218
+ N = B if cu_seqlens is None else len(cu_seqlens) - 1
219
+ BK, BV = triton.next_power_of_2(K), min(triton.next_power_of_2(V), 32)
220
+ NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV)
221
+ assert NK == 1, "NK > 1 is not supported yet"
222
+ num_stages = 3
223
+ num_warps = 1
224
+
225
+ o = q.new_empty(NK, *v.shape)
226
+ if inplace_final_state:
227
+ final_state = initial_state
228
+ else:
229
+ final_state = q.new_empty(T, HV, V, K, dtype=initial_state.dtype)
230
+
231
+ stride_init_state_token = initial_state.stride(0)
232
+ stride_final_state_token = final_state.stride(0)
233
+
234
+ if ssm_state_indices is None:
235
+ stride_indices_seq, stride_indices_tok = 1, 1
236
+ elif ssm_state_indices.ndim == 1:
237
+ stride_indices_seq, stride_indices_tok = ssm_state_indices.stride(0), 1
238
+ else:
239
+ stride_indices_seq, stride_indices_tok = ssm_state_indices.stride()
240
+
241
+ grid = (NK, NV, N * HV)
242
+ fused_recurrent_gated_delta_rule_fwd_kernel[grid](
243
+ q=q,
244
+ k=k,
245
+ v=v,
246
+ g=g,
247
+ beta=beta,
248
+ o=o,
249
+ h0=initial_state,
250
+ ht=final_state,
251
+ cu_seqlens=cu_seqlens,
252
+ ssm_state_indices=ssm_state_indices,
253
+ num_accepted_tokens=num_accepted_tokens,
254
+ scale=scale,
255
+ N=N,
256
+ T=T,
257
+ B=B,
258
+ H=H,
259
+ HV=HV,
260
+ K=K,
261
+ V=V,
262
+ BK=BK,
263
+ BV=BV,
264
+ stride_init_state_token=stride_init_state_token,
265
+ stride_final_state_token=stride_final_state_token,
266
+ stride_indices_seq=stride_indices_seq,
267
+ stride_indices_tok=stride_indices_tok,
268
+ IS_BETA_HEADWISE=beta.ndim == v.ndim,
269
+ USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,
270
+ INPLACE_FINAL_STATE=inplace_final_state,
271
+ IS_KDA=False,
272
+ SIGMOID_BETA=False,
273
+ a_log=None,
274
+ g_bias=None,
275
+ COMPUTE_GATE=False,
276
+ SAFE_GATE=True,
277
+ LOWER_BOUND=-5.0,
278
+ num_warps=num_warps,
279
+ num_stages=num_stages,
280
+ )
281
+ o = o.squeeze(0)
282
+ return o, final_state
283
+
284
+
285
+@triton.jit
286
+def fused_recurrent_gated_delta_rule_packed_decode_kernel(
287
+ mixed_qkv,
288
+ a,
289
+ b,
290
+ A_log,
291
+ dt_bias,
292
+ o,
293
+ h0,
294
+ ht,
295
+ ssm_state_indices,
296
+ scale,
297
+ stride_mixed_qkv_tok: tl.constexpr,
298
+ stride_a_tok: tl.constexpr,
299
+ stride_b_tok: tl.constexpr,
300
+ stride_init_state_token: tl.constexpr,
301
+ stride_final_state_token: tl.constexpr,
302
+ stride_indices_seq: tl.constexpr,
303
+ H: tl.constexpr,
304
+ HV: tl.constexpr,
305
+ K: tl.constexpr,
306
+ V: tl.constexpr,
307
+ BK: tl.constexpr,
308
+ BV: tl.constexpr,
309
+ SOFTPLUS_THRESHOLD: tl.constexpr,
310
+ USE_QK_L2NORM_IN_KERNEL: tl.constexpr,
311
+ SPLIT_BATCH_HEAD_GRID: tl.constexpr,
312
+):
313
+ if SPLIT_BATCH_HEAD_GRID:
314
+ i_v, i_hv, i_n = tl.program_id(0), tl.program_id(1), tl.program_id(2)
315
+ else:
316
+ i_v, i_nh = tl.program_id(0), tl.program_id(1)
317
+ i_n, i_hv = i_nh // HV, i_nh % HV
318
+ i_h = i_hv // (HV // H)
319
+
320
+ o_k = tl.arange(0, BK)
321
+ o_v = i_v * BV + tl.arange(0, BV)
322
+ mask_k = o_k < K
323
+ mask_v = o_v < V
324
+ mask_h = mask_v[:, None] & mask_k[None, :]
325
+
326
+ state_idx = tl.load(ssm_state_indices + i_n * stride_indices_seq).to(tl.int64)
327
+ p_o = o + (i_n * HV + i_hv) * V + o_v
328
+
329
+ # Skip if state index is invalid (NULL_BLOCK_ID=0)
330
+ if state_idx <= 0:
331
+ zero = tl.zeros([BV], dtype=tl.float32).to(p_o.dtype.element_ty)
332
+ tl.store(p_o, zero, mask=mask_v)
333
+ return
334
+
335
+ p_h0 = h0 + state_idx * stride_init_state_token
336
+ p_h0 = p_h0 + i_hv * V * K + o_v[:, None] * K + o_k[None, :]
337
+ b_h = tl.load(p_h0, mask=mask_h, other=0).to(tl.float32)
338
+
339
+ p_mixed = mixed_qkv + i_n * stride_mixed_qkv_tok
340
+ q_off = i_h * K + o_k
341
+ k_off = (H * K) + i_h * K + o_k
342
+ v_off = (2 * H * K) + i_hv * V + o_v
343
+ b_q = tl.load(p_mixed + q_off, mask=mask_k, other=0).to(tl.float32)
344
+ b_k = tl.load(p_mixed + k_off, mask=mask_k, other=0).to(tl.float32)
345
+ b_v = tl.load(p_mixed + v_off, mask=mask_v, other=0).to(tl.float32)
346
+
347
+ if USE_QK_L2NORM_IN_KERNEL:
348
+ b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6)
349
+ b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6)
350
+ b_q = b_q * scale
351
+
352
+ a_val = tl.load(a + i_n * stride_a_tok + i_hv).to(tl.float32)
353
+ b_val = tl.load(b + i_n * stride_b_tok + i_hv).to(tl.float32)
354
+ A_log_val = tl.load(A_log + i_hv).to(tl.float32)
355
+ dt_bias_val = tl.load(dt_bias + i_hv).to(tl.float32)
356
+ x = a_val + dt_bias_val
357
+ softplus_x = tl.where(x <= SOFTPLUS_THRESHOLD, tl.log(1.0 + tl.exp(x)), x)
358
+ g_val = -tl.exp(A_log_val) * softplus_x
359
+ beta_val = tl.sigmoid(b_val)
360
+
361
+ b_h *= exp(g_val)
362
+ b_v -= tl.sum(b_h * b_k[None, :], 1)
363
+ b_v *= beta_val
364
+ b_h += b_v[:, None] * b_k[None, :]
365
+ b_o = tl.sum(b_h * b_q[None, :], 1)
366
+ tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v)
367
+
368
+ p_ht = ht + state_idx * stride_final_state_token
369
+ p_ht = p_ht + i_hv * V * K + o_v[:, None] * K + o_k[None, :]
370
+ tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h)
371
+
372
+
373
+def fused_recurrent_gated_delta_rule_packed_decode(
374
+ mixed_qkv: torch.Tensor,
375
+ a: torch.Tensor,
376
+ b: torch.Tensor,
377
+ A_log: torch.Tensor,
378
+ dt_bias: torch.Tensor,
379
+ scale: float,
380
+ initial_state: torch.Tensor,
381
+ out: torch.Tensor,
382
+ ssm_state_indices: torch.Tensor,
383
+ use_qk_l2norm_in_kernel: bool = False,
384
+) -> tuple[torch.Tensor, torch.Tensor]:
385
+ if mixed_qkv.ndim != 2:
386
+ raise ValueError(
387
+ f"`mixed_qkv` must be a 2D tensor (got ndim={mixed_qkv.ndim})."
388
+ )
389
+ if mixed_qkv.stride(-1) != 1:
390
+ raise ValueError("`mixed_qkv` must be contiguous in the last dim.")
391
+ if a.ndim != 2 or b.ndim != 2:
392
+ raise ValueError(
393
+ f"`a` and `b` must be 2D tensors (got a.ndim={a.ndim}, b.ndim={b.ndim})."
394
+ )
395
+ if a.stride(-1) != 1 or b.stride(-1) != 1:
396
+ raise ValueError("`a`/`b` must be contiguous in the last dim.")
397
+ if A_log.ndim != 1 or dt_bias.ndim != 1:
398
+ raise ValueError("`A_log`/`dt_bias` must be 1D tensors.")
399
+ if A_log.stride(0) != 1 or dt_bias.stride(0) != 1:
400
+ raise ValueError("`A_log`/`dt_bias` must be contiguous.")
401
+ if ssm_state_indices.ndim != 1:
402
+ raise ValueError(
403
+ f"`ssm_state_indices` must be 1D for packed decode (got ndim={ssm_state_indices.ndim})."
404
+ )
405
+ if not out.is_contiguous():
406
+ raise ValueError("`out` must be contiguous.")
407
+
408
+ dev = mixed_qkv.device
409
+ if (
410
+ a.device != dev
411
+ or b.device != dev
412
+ or A_log.device != dev
413
+ or dt_bias.device != dev
414
+ or initial_state.device != dev
415
+ or out.device != dev
416
+ or ssm_state_indices.device != dev
417
+ ):
418
+ raise ValueError("All inputs must be on the same device.")
419
+
420
+ B = mixed_qkv.shape[0]
421
+ if a.shape[0] != B or b.shape[0] != B:
422
+ raise ValueError(
423
+ "Mismatched batch sizes: "
424
+ f"mixed_qkv.shape[0]={B}, a.shape[0]={a.shape[0]}, b.shape[0]={b.shape[0]}."
425
+ )
426
+ if ssm_state_indices.shape[0] != B:
427
+ raise ValueError(
428
+ f"`ssm_state_indices` must have shape [B] (got {tuple(ssm_state_indices.shape)}; expected ({B},))."
429
+ )
430
+
431
+ if initial_state.ndim != 4:
432
+ raise ValueError(
433
+ f"`initial_state` must be a 4D tensor (got ndim={initial_state.ndim})."
434
+ )
435
+ if initial_state.stride(-1) != 1:
436
+ raise ValueError("`initial_state` must be contiguous in the last dim.")
437
+ HV, V, K = initial_state.shape[-3:]
438
+ if a.shape[1] != HV or b.shape[1] != HV:
439
+ raise ValueError(
440
+ f"`a`/`b` must have shape [B, HV] with HV={HV} (got a.shape={tuple(a.shape)}, b.shape={tuple(b.shape)})."
441
+ )
442
+ if A_log.numel() != HV or dt_bias.numel() != HV:
443
+ raise ValueError(
444
+ f"`A_log` and `dt_bias` must have {HV} elements (got A_log.numel()={A_log.numel()}, dt_bias.numel()={dt_bias.numel()})."
445
+ )
446
+ if out.shape != (B, 1, HV, V):
447
+ raise ValueError(
448
+ f"`out` must have shape {(B, 1, HV, V)} (got out.shape={tuple(out.shape)})."
449
+ )
450
+
451
+ qkv_dim = mixed_qkv.shape[1]
452
+ qk_dim = qkv_dim - HV * V
453
+ if qk_dim <= 0 or qk_dim % 2 != 0:
454
+ raise ValueError(
455
+ f"Invalid packed `mixed_qkv` last dim={qkv_dim} for HV={HV}, V={V}."
456
+ )
457
+ q_dim = qk_dim // 2
458
+ if q_dim % K != 0:
459
+ raise ValueError(f"Invalid packed Q size {q_dim}: must be divisible by K={K}.")
460
+ H = q_dim // K
461
+ if H <= 0 or HV % H != 0:
462
+ raise ValueError(
463
+ f"Invalid head config inferred from mixed_qkv: H={H}, HV={HV}."
464
+ )
465
+
466
+ BK = triton.next_power_of_2(K)
467
+ if triton.cdiv(K, BK) != 1:
468
+ raise ValueError(
469
+ f"Packed decode kernel only supports NK=1 (got K={K}, BK={BK})."
470
+ )
471
+ BV = min(triton.next_power_of_2(V), 32)
472
+ num_stages = 3
473
+ num_warps = 1
474
+
475
+ stride_mixed_qkv_tok = mixed_qkv.stride(0)
476
+ stride_a_tok = a.stride(0)
477
+ stride_b_tok = b.stride(0)
478
+ stride_init_state_token = initial_state.stride(0)
479
+ stride_final_state_token = initial_state.stride(0)
480
+ stride_indices_seq = ssm_state_indices.stride(0)
481
+
482
+ NV = triton.cdiv(V, BV)
483
+ # CUDA limits grid Y/Z dimensions to 65535.
484
+ split_batch_head_grid = B * HV > 65535
485
+ grid = (NV, HV, B) if split_batch_head_grid else (NV, B * HV)
486
+ fused_recurrent_gated_delta_rule_packed_decode_kernel[grid](
487
+ mixed_qkv=mixed_qkv,
488
+ a=a,
489
+ b=b,
490
+ A_log=A_log,
491
+ dt_bias=dt_bias,
492
+ o=out,
493
+ h0=initial_state,
494
+ ht=initial_state,
495
+ ssm_state_indices=ssm_state_indices,
496
+ scale=scale,
497
+ stride_mixed_qkv_tok=stride_mixed_qkv_tok,
498
+ stride_a_tok=stride_a_tok,
499
+ stride_b_tok=stride_b_tok,
500
+ stride_init_state_token=stride_init_state_token,
501
+ stride_final_state_token=stride_final_state_token,
502
+ stride_indices_seq=stride_indices_seq,
503
+ H=H,
504
+ HV=HV,
505
+ K=K,
506
+ V=V,
507
+ BK=BK,
508
+ BV=BV,
509
+ SOFTPLUS_THRESHOLD=20.0,
510
+ USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,
511
+ SPLIT_BATCH_HEAD_GRID=split_batch_head_grid,
512
+ num_warps=num_warps,
513
+ num_stages=num_stages,
514
+ )
515
+ return out, initial_state
516
+
517
+
518
+class FusedRecurrentFunction(torch.autograd.Function):
519
+ @staticmethod
520
+ def forward(
521
+ ctx,
522
+ q: torch.Tensor,
523
+ k: torch.Tensor,
524
+ v: torch.Tensor,
525
+ g: torch.Tensor,
526
+ beta: torch.Tensor,
527
+ scale: float,
528
+ initial_state: torch.Tensor,
529
+ inplace_final_state: bool = True,
530
+ cu_seqlens: torch.Tensor | None = None,
531
+ ssm_state_indices: torch.Tensor | None = None,
532
+ num_accepted_tokens: torch.Tensor | None = None,
533
+ use_qk_l2norm_in_kernel: bool = False,
534
+ ):
535
+ o, final_state = fused_recurrent_gated_delta_rule_fwd(
536
+ q=q.contiguous(),
537
+ k=k.contiguous(),
538
+ v=v.contiguous(),
539
+ g=g.contiguous(),
540
+ beta=beta.contiguous(),
541
+ scale=scale,
542
+ initial_state=initial_state,
543
+ inplace_final_state=inplace_final_state,
544
+ cu_seqlens=cu_seqlens,
545
+ ssm_state_indices=ssm_state_indices,
546
+ num_accepted_tokens=num_accepted_tokens,
547
+ use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
548
+ )
549
+
550
+ return o, final_state
551
+
552
+
553
+def fused_recurrent_gated_delta_rule(
554
+ q: torch.Tensor,
555
+ k: torch.Tensor,
556
+ v: torch.Tensor,
557
+ g: torch.Tensor,
558
+ beta: torch.Tensor = None,
559
+ scale: float = None,
560
+ initial_state: torch.Tensor = None,
561
+ inplace_final_state: bool = True,
562
+ cu_seqlens: torch.Tensor | None = None,
563
+ ssm_state_indices: torch.Tensor | None = None,
564
+ num_accepted_tokens: torch.Tensor | None = None,
565
+ use_qk_l2norm_in_kernel: bool = False,
566
+) -> tuple[torch.Tensor, torch.Tensor]:
567
+ r"""
568
+ Args:
569
+ q (torch.Tensor):
570
+ queries of shape `[B, T, H, K]`.
571
+ k (torch.Tensor):
572
+ keys of shape `[B, T, H, K]`.
573
+ v (torch.Tensor):
574
+ values of shape `[B, T, HV, V]`.
575
+ GVA is applied if `HV > H`.
576
+ g (torch.Tensor):
577
+ g (decays) of shape `[B, T, HV]`.
578
+ beta (torch.Tensor):
579
+ betas of shape `[B, T, HV]`.
580
+ scale (Optional[int]):
581
+ Scale factor for the RetNet attention scores.
582
+ If not provided, it will default to `1 / sqrt(K)`. Default: `None`.
583
+ initial_state (Optional[torch.Tensor]):
584
+ Initial state of shape `[N, HV, V, K]` for `N` input sequences.
585
+ For equal-length input sequences, `N` equals the batch size `B`.
586
+ Default: `None`.
587
+ inplace_final_state: bool:
588
+ Whether to store the final state in-place to save memory.
589
+ Default: `True`.
590
+ cu_seqlens (torch.Tensor):
591
+ Cumulative sequence lengths of shape `[N+1]` used for variable-length training,
592
+ consistent with the FlashAttention API.
593
+ ssm_state_indices (Optional[torch.Tensor]):
594
+ Indices to map the input sequences to the initial/final states.
595
+ num_accepted_tokens (Optional[torch.Tensor]):
596
+ Number of accepted tokens for each sequence during decoding.
597
+
598
+ Returns:
599
+ o (torch.Tensor):
600
+ Outputs of shape `[B, T, HV, V]`.
601
+ final_state (torch.Tensor):
602
+ Final state of shape `[N, HV, V, K]`.
603
+
604
+ Examples::
605
+ >>> import torch
606
+ >>> import torch.nn.functional as F
607
+ >>> from einops import rearrange
608
+ >>> from fla.ops.gated_delta_rule import fused_recurrent_gated_delta_rule
609
+ # inputs with equal lengths
610
+ >>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512
611
+ >>> q = torch.randn(B, T, H, K, device='cuda')
612
+ >>> k = F.normalize(torch.randn(B, T, H, K, device='cuda'), p=2, dim=-1)
613
+ >>> v = torch.randn(B, T, HV, V, device='cuda')
614
+ >>> g = F.logsigmoid(torch.rand(B, T, HV, device='cuda'))
615
+ >>> beta = torch.rand(B, T, HV, device='cuda').sigmoid()
616
+ >>> h0 = torch.randn(B, HV, V, K, device='cuda')
617
+ >>> o, ht = fused_gated_recurrent_delta_rule(
618
+ q, k, v, g, beta,
619
+ initial_state=h0,
620
+ )
621
+ # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required
622
+ >>> q, k, v, g, beta = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, g, beta))
623
+ # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected
624
+ >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.int32)
625
+ >>> o_var, ht_var = fused_gated_recurrent_delta_rule(
626
+ q, k, v, g, beta,
627
+ initial_state=h0,
628
+ cu_seqlens=cu_seqlens
629
+ )
630
+ """
631
+ if cu_seqlens is not None and q.shape[0] != 1:
632
+ raise ValueError(
633
+ f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`."
634
+ f"Please flatten variable-length inputs before processing."
635
+ )
636
+ if scale is None:
637
+ scale = k.shape[-1] ** -0.5
638
+ else:
639
+ assert scale > 0, "scale must be positive"
640
+ if beta is None:
641
+ beta = torch.ones_like(q[..., 0])
642
+ o, final_state = FusedRecurrentFunction.apply(
643
+ q,
644
+ k,
645
+ v,
646
+ g,
647
+ beta,
648
+ scale,
649
+ initial_state,
650
+ inplace_final_state,
651
+ cu_seqlens,
652
+ ssm_state_indices,
653
+ num_accepted_tokens,
654
+ use_qk_l2norm_in_kernel,
655
+ )
656
+ return o, final_state
vllm/models/glm5next/amd/ops/third_party/kda/kernels.py
new
+1362
@@ -0,0 +1,1362 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang
4
+# mypy: ignore-errors
5
+#
6
+# This file contains code copied from the flash-linear-attention project.
7
+# The original source code was licensed under the MIT license and included
8
+# the following copyright notice:
9
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
10
+# ruff: noqa: E501
11
+
12
+
13
+import torch
14
+
15
+from vllm.third_party.flash_linear_attention.ops.chunk_delta_h import (
16
+ chunk_gated_delta_rule_fwd_h,
17
+)
18
+from vllm.third_party.flash_linear_attention.ops.cumsum import chunk_local_cumsum
19
+from vllm.third_party.flash_linear_attention.ops.index import prepare_chunk_indices
20
+from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd
21
+from vllm.third_party.flash_linear_attention.ops.op import exp2, log
22
+from vllm.third_party.flash_linear_attention.ops.solve_tril import solve_tril
23
+from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE, is_amd
24
+from vllm.triton_utils import tl, triton
25
+from vllm.utils.math_utils import RCP_LN2, cdiv, next_power_of_2
26
+
27
+from .fused_recurrent import fused_recurrent_gated_delta_rule_fwd_kernel
28
+
29
+BT_LIST_AUTOTUNE = [32, 64, 128]
30
+NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if is_amd else [4, 8, 16, 32]
31
+
32
+
33
+def fused_recurrent_kda_fwd(
34
+ q: torch.Tensor,
35
+ k: torch.Tensor,
36
+ v: torch.Tensor,
37
+ g: torch.Tensor,
38
+ beta: torch.Tensor,
39
+ scale: float,
40
+ initial_state: torch.Tensor,
41
+ inplace_final_state: bool = True,
42
+ cu_seqlens: torch.Tensor | None = None,
43
+ ssm_state_indices: torch.Tensor | None = None,
44
+ num_accepted_tokens: torch.Tensor | None = None,
45
+ use_qk_l2norm_in_kernel: bool = False,
46
+ out: torch.Tensor | None = None,
47
+ sigmoid_beta: bool = False,
48
+ a_log: torch.Tensor | None = None,
49
+ g_bias: torch.Tensor | None = None,
50
+ compute_gate: bool = False,
51
+ lower_bound: float | None = -5.0,
52
+) -> tuple[torch.Tensor, torch.Tensor]:
53
+ B, T, H, K, V = *k.shape, v.shape[-1]
54
+ HV = v.shape[2]
55
+ N = B if cu_seqlens is None else len(cu_seqlens) - 1
56
+ BK, BV = next_power_of_2(K), min(next_power_of_2(V), 8)
57
+ NK, NV = cdiv(K, BK), cdiv(V, BV)
58
+ assert NK == 1, "NK > 1 is not supported yet"
59
+ num_stages = 3
60
+ num_warps = 1
61
+
62
+ if compute_gate:
63
+ assert a_log is not None and g_bias is not None, (
64
+ "compute_gate requires a_log and g_bias"
65
+ )
66
+ assert lower_bound is not None, (
67
+ "compute_gate implements the bounded (safe_gate) branch only"
68
+ )
69
+ a_log = a_log.reshape(-1).contiguous()
70
+ g_bias = g_bias.reshape(-1).contiguous()
71
+
72
+ if out is None:
73
+ o = torch.empty_like(k)
74
+ else:
75
+ # Caller-provided output buffer; must be layout-compatible with the
76
+ # tensor the kernel indexes (contiguous, same shape/dtype as k).
77
+ assert out.shape == k.shape and out.dtype == k.dtype
78
+ assert out.is_contiguous()
79
+ o = out
80
+ if inplace_final_state:
81
+ final_state = initial_state
82
+ else:
83
+ final_state = q.new_empty(T, HV, V, K, dtype=initial_state.dtype)
84
+
85
+ stride_init_state_token = initial_state.stride(0)
86
+ stride_final_state_token = final_state.stride(0)
87
+
88
+ if ssm_state_indices is None:
89
+ stride_indices_seq, stride_indices_tok = 1, 1
90
+ elif ssm_state_indices.ndim == 1:
91
+ stride_indices_seq, stride_indices_tok = ssm_state_indices.stride(0), 1
92
+ else:
93
+ stride_indices_seq, stride_indices_tok = ssm_state_indices.stride()
94
+
95
+ grid = (NK, NV, N * HV)
96
+ fused_recurrent_gated_delta_rule_fwd_kernel[grid](
97
+ q=q,
98
+ k=k,
99
+ v=v,
100
+ g=g,
101
+ beta=beta,
102
+ o=o,
103
+ h0=initial_state,
104
+ ht=final_state,
105
+ cu_seqlens=cu_seqlens,
106
+ ssm_state_indices=ssm_state_indices,
107
+ num_accepted_tokens=num_accepted_tokens,
108
+ scale=scale,
109
+ N=N,
110
+ T=T,
111
+ B=B,
112
+ H=H,
113
+ HV=HV,
114
+ K=K,
115
+ V=V,
116
+ BK=BK,
117
+ BV=BV,
118
+ stride_init_state_token=stride_init_state_token,
119
+ stride_final_state_token=stride_final_state_token,
120
+ stride_indices_seq=stride_indices_seq,
121
+ stride_indices_tok=stride_indices_tok,
122
+ IS_BETA_HEADWISE=beta.ndim == v.ndim,
123
+ USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,
124
+ INPLACE_FINAL_STATE=inplace_final_state,
125
+ IS_KDA=True,
126
+ SIGMOID_BETA=sigmoid_beta,
127
+ a_log=a_log,
128
+ g_bias=g_bias,
129
+ COMPUTE_GATE=compute_gate,
130
+ SAFE_GATE=True,
131
+ LOWER_BOUND=lower_bound if lower_bound is not None else -5.0,
132
+ num_warps=num_warps,
133
+ num_stages=num_stages,
134
+ )
135
+
136
+ return o, final_state
137
+
138
+
139
+def fused_recurrent_kda(
140
+ q: torch.Tensor,
141
+ k: torch.Tensor,
142
+ v: torch.Tensor,
143
+ g: torch.Tensor,
144
+ beta: torch.Tensor = None,
145
+ scale: float = None,
146
+ initial_state: torch.Tensor = None,
147
+ inplace_final_state: bool = True,
148
+ use_qk_l2norm_in_kernel: bool = True,
149
+ cu_seqlens: torch.Tensor | None = None,
150
+ ssm_state_indices: torch.LongTensor | None = None,
151
+ num_accepted_tokens: torch.Tensor | None = None,
152
+ out: torch.Tensor | None = None,
153
+ sigmoid_beta: bool = False,
154
+ a_log: torch.Tensor | None = None,
155
+ g_bias: torch.Tensor | None = None,
156
+ compute_gate: bool = False,
157
+ lower_bound: float | None = -5.0,
158
+ **kwargs,
159
+) -> tuple[torch.Tensor, torch.Tensor]:
160
+ if cu_seqlens is not None and q.shape[0] != 1:
161
+ raise ValueError(
162
+ f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`."
163
+ f"Please flatten variable-length inputs before processing."
164
+ )
165
+ if scale is None:
166
+ scale = k.shape[-1] ** -0.5
167
+
168
+ o, final_state = fused_recurrent_kda_fwd(
169
+ q=q.contiguous(),
170
+ k=k.contiguous(),
171
+ v=v.contiguous(),
172
+ g=g.contiguous(),
173
+ beta=beta.contiguous(),
174
+ scale=scale,
175
+ initial_state=initial_state,
176
+ inplace_final_state=inplace_final_state,
177
+ cu_seqlens=cu_seqlens,
178
+ ssm_state_indices=ssm_state_indices,
179
+ num_accepted_tokens=num_accepted_tokens,
180
+ use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
181
+ out=out,
182
+ sigmoid_beta=sigmoid_beta,
183
+ a_log=a_log,
184
+ g_bias=g_bias,
185
+ compute_gate=compute_gate,
186
+ lower_bound=lower_bound,
187
+ )
188
+ return o, final_state
189
+
190
+
191
+@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
192
+@triton.autotune(
193
+ configs=[
194
+ triton.Config({"BK": BK}, num_warps=num_warps, num_stages=num_stages)
195
+ for BK in [32, 64]
196
+ for num_warps in [1, 2, 4, 8]
197
+ for num_stages in [2, 3, 4]
198
+ ],
199
+ key=["BC"],
200
+)
201
+@triton.jit(do_not_specialize=["T"])
202
+def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter(
203
+ q,
204
+ k,
205
+ g,
206
+ beta,
207
+ A,
208
+ Aqk,
209
+ scale,
210
+ cu_seqlens,
211
+ chunk_indices,
212
+ T,
213
+ H: tl.constexpr,
214
+ K: tl.constexpr,
215
+ BT: tl.constexpr,
216
+ BC: tl.constexpr,
217
+ BK: tl.constexpr,
218
+ NC: tl.constexpr,
219
+ IS_VARLEN: tl.constexpr,
220
+):
221
+ i_t, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
222
+ i_b, i_h = i_bh // H, i_bh % H
223
+ i_i, i_j = i_c // NC, i_c % NC
224
+ if IS_VARLEN:
225
+ i_n, i_t = (
226
+ tl.load(chunk_indices + i_t * 2).to(tl.int32),
227
+ tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
228
+ )
229
+ bos, eos = (
230
+ tl.load(cu_seqlens + i_n).to(tl.int32),
231
+ tl.load(cu_seqlens + i_n + 1).to(tl.int32),
232
+ )
233
+ T = eos - bos
234
+ else:
235
+ bos, eos = i_b * T, i_b * T + T
236
+
237
+ if i_t * BT + i_i * BC >= T:
238
+ return
239
+ if i_i <= i_j:
240
+ return
241
+
242
+ q += (bos * H + i_h) * K
243
+ k += (bos * H + i_h) * K
244
+ g += (bos * H + i_h) * K
245
+ A += (bos * H + i_h) * BT
246
+ Aqk += (bos * H + i_h) * BT
247
+
248
+ p_b = tl.make_block_ptr(
249
+ beta + bos * H + i_h, (T,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,)
250
+ )
251
+ b_b = tl.load(p_b, boundary_check=(0,))
252
+
253
+ b_A = tl.zeros([BC, BC], dtype=tl.float32)
254
+ b_Aqk = tl.zeros([BC, BC], dtype=tl.float32)
255
+ for i_k in range(tl.cdiv(K, BK)):
256
+ p_q = tl.make_block_ptr(
257
+ q, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)
258
+ )
259
+ p_k = tl.make_block_ptr(
260
+ k, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)
261
+ )
262
+ p_g = tl.make_block_ptr(
263
+ g, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)
264
+ )
265
+ b_kt = tl.make_block_ptr(
266
+ k, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)
267
+ )
268
+ p_gk = tl.make_block_ptr(
269
+ g, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)
270
+ )
271
+
272
+ o_k = i_k * BK + tl.arange(0, BK)
273
+ m_k = o_k < K
274
+ # [BK,]
275
+ b_gn = tl.load(g + (i_t * BT + i_i * BC) * H * K + o_k, mask=m_k, other=0)
276
+ # [BC, BK]
277
+ b_g = tl.load(p_g, boundary_check=(0, 1))
278
+ b_k = tl.load(p_k, boundary_check=(0, 1)) * exp2(b_g - b_gn[None, :])
279
+ # [BK, BC]
280
+ b_gk = tl.load(p_gk, boundary_check=(0, 1))
281
+ b_kt = tl.load(b_kt, boundary_check=(0, 1))
282
+ # [BC, BC]
283
+ b_ktg = b_kt * exp2(b_gn[:, None] - b_gk)
284
+ b_A += tl.dot(b_k, b_ktg)
285
+
286
+ b_q = tl.load(p_q, boundary_check=(0, 1))
287
+ b_qg = b_q * exp2(b_g - b_gn[None, :]) * scale
288
+ b_Aqk += tl.dot(b_qg, b_ktg)
289
+
290
+ b_A *= b_b[:, None]
291
+
292
+ p_A = tl.make_block_ptr(
293
+ A, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)
294
+ )
295
+ tl.store(p_A, b_A.to(A.dtype.element_ty), boundary_check=(0, 1))
296
+ p_Aqk = tl.make_block_ptr(
297
+ Aqk, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)
298
+ )
299
+ tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1))
300
+
301
+
302
+@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
303
+@triton.autotune(
304
+ configs=[triton.Config({}, num_warps=num_warps) for num_warps in [1, 2, 4, 8]],
305
+ key=["BK", "BT"],
306
+)
307
+@triton.jit(do_not_specialize=["T"])
308
+def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra(
309
+ q,
310
+ k,
311
+ g,
312
+ beta,
313
+ A,
314
+ Aqk,
315
+ scale,
316
+ cu_seqlens,
317
+ chunk_indices,
318
+ T,
319
+ H: tl.constexpr,
320
+ K: tl.constexpr,
321
+ BT: tl.constexpr,
322
+ BC: tl.constexpr,
323
+ BK: tl.constexpr,
324
+ IS_VARLEN: tl.constexpr,
325
+):
326
+ i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
327
+ i_b, i_h = i_bh // H, i_bh % H
328
+ if IS_VARLEN:
329
+ i_n, i_t = (
330
+ tl.load(chunk_indices + i_t * 2).to(tl.int32),
331
+ tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
332
+ )
333
+ bos, eos = (
334
+ tl.load(cu_seqlens + i_n).to(tl.int32),
335
+ tl.load(cu_seqlens + i_n + 1).to(tl.int32),
336
+ )
337
+ T = eos - bos
338
+ else:
339
+ bos, eos = i_b * T, i_b * T + T
340
+
341
+ if i_t * BT + i_i * BC >= T:
342
+ return
343
+
344
+ o_i = tl.arange(0, BC)
345
+ o_k = tl.arange(0, BK)
346
+ m_k = o_k < K
347
+ m_A = (i_t * BT + i_i * BC + o_i) < T
348
+ o_A = (bos + i_t * BT + i_i * BC + o_i) * H * BT + i_h * BT + i_i * BC
349
+
350
+ p_q = tl.make_block_ptr(
351
+ q + (bos * H + i_h) * K,
352
+ (T, K),
353
+ (H * K, 1),
354
+ (i_t * BT + i_i * BC, 0),
355
+ (BC, BK),
356
+ (1, 0),
357
+ )
358
+ p_k = tl.make_block_ptr(
359
+ k + (bos * H + i_h) * K,
360
+ (T, K),
361
+ (H * K, 1),
362
+ (i_t * BT + i_i * BC, 0),
363
+ (BC, BK),
364
+ (1, 0),
365
+ )
366
+ p_g = tl.make_block_ptr(
367
+ g + (bos * H + i_h) * K,
368
+ (T, K),
369
+ (H * K, 1),
370
+ (i_t * BT + i_i * BC, 0),
371
+ (BC, BK),
372
+ (1, 0),
373
+ )
374
+ b_q = tl.load(p_q, boundary_check=(0, 1))
375
+ b_k = tl.load(p_k, boundary_check=(0, 1))
376
+ b_g = tl.load(p_g, boundary_check=(0, 1))
377
+
378
+ p_b = beta + (bos + i_t * BT + i_i * BC + o_i) * H + i_h
379
+ b_k = b_k * tl.load(p_b, mask=m_A, other=0)[:, None]
380
+
381
+ p_kt = k + (bos + i_t * BT + i_i * BC) * H * K + i_h * K + o_k
382
+ p_gk = g + (bos + i_t * BT + i_i * BC) * H * K + i_h * K + o_k
383
+
384
+ for j in range(0, min(BC, T - i_t * BT - i_i * BC)):
385
+ b_kt = tl.load(p_kt, mask=m_k, other=0).to(tl.float32)
386
+ b_gk = tl.load(p_gk, mask=m_k, other=0).to(tl.float32)
387
+ b_ktg = b_kt[None, :] * exp2(b_g - b_gk[None, :])
388
+ b_A = tl.sum(b_k * b_ktg, 1)
389
+ b_A = tl.where(o_i > j, b_A, 0.0)
390
+ b_Aqk = tl.sum(b_q * b_ktg, 1)
391
+ b_Aqk = tl.where(o_i >= j, b_Aqk * scale, 0.0)
392
+ tl.store(A + o_A + j, b_A, mask=m_A)
393
+ tl.store(Aqk + o_A + j, b_Aqk, mask=m_A)
394
+ p_kt += H * K
395
+ p_gk += H * K
396
+
397
+
398
+def chunk_kda_scaled_dot_kkt_fwd(
399
+ q: torch.Tensor,
400
+ k: torch.Tensor,
401
+ gk: torch.Tensor | None = None,
402
+ beta: torch.Tensor | None = None,
403
+ scale: float | None = None,
404
+ cu_seqlens: torch.Tensor | None = None,
405
+ chunk_indices: torch.Tensor | None = None,
406
+ chunk_size: int = FLA_CHUNK_SIZE,
407
+ output_dtype: torch.dtype = torch.float32,
408
+) -> tuple[torch.Tensor, torch.Tensor]:
409
+ r"""
410
+ Compute beta * K * K^T.
411
+
412
+ Args:
413
+ k (torch.Tensor):
414
+ The key tensor of shape `[B, T, H, K]`.
415
+ beta (torch.Tensor):
416
+ The beta tensor of shape `[B, T, H]`.
417
+ gk (torch.Tensor):
418
+ The cumulative sum of the gate tensor of shape `[B, T, H, K]` applied to the key tensor. Default: `None`.
419
+ cu_seqlens (torch.Tensor):
420
+ The cumulative sequence lengths of the input tensor.
421
+ Default: None
422
+ chunk_size (int):
423
+ The chunk size. Default: 64.
424
+ output_dtype (torch.dtype):
425
+ The dtype of the output tensor. Default: `torch.float32`
426
+
427
+ Returns:
428
+ beta * K * K^T of shape `[B, T, H, BT]` where `BT` is the chunk size.
429
+ """
430
+ B, T, H, K = k.shape
431
+ assert K <= 256
432
+ BT = chunk_size
433
+ if chunk_indices is None and cu_seqlens is not None:
434
+ chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
435
+ NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
436
+
437
+ BC = min(16, BT)
438
+ NC = cdiv(BT, BC)
439
+ BK = max(next_power_of_2(K), 16)
440
+ A = torch.zeros(B, T, H, BT, device=k.device, dtype=output_dtype)
441
+ Aqk = torch.zeros(B, T, H, BT, device=k.device, dtype=output_dtype)
442
+ grid = (NT, NC * NC, B * H)
443
+ chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter[grid](
444
+ q=q,
445
+ k=k,
446
+ g=gk,
447
+ beta=beta,
448
+ A=A,
449
+ Aqk=Aqk,
450
+ scale=scale,
451
+ cu_seqlens=cu_seqlens,
452
+ chunk_indices=chunk_indices,
453
+ T=T,
454
+ H=H,
455
+ K=K,
456
+ BT=BT,
457
+ BC=BC,
458
+ NC=NC,
459
+ )
460
+
461
+ grid = (NT, NC, B * H)
462
+ chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra[grid](
463
+ q=q,
464
+ k=k,
465
+ g=gk,
466
+ beta=beta,
467
+ A=A,
468
+ Aqk=Aqk,
469
+ scale=scale,
470
+ cu_seqlens=cu_seqlens,
471
+ chunk_indices=chunk_indices,
472
+ T=T,
473
+ H=H,
474
+ K=K,
475
+ BT=BT,
476
+ BC=BC,
477
+ BK=BK,
478
+ )
479
+ return A, Aqk
480
+
481
+
482
+@triton.heuristics(
483
+ {
484
+ "STORE_QG": lambda args: args["qg"] is not None,
485
+ "STORE_KG": lambda args: args["kg"] is not None,
486
+ "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
487
+ }
488
+)
489
+@triton.autotune(
490
+ configs=[
491
+ triton.Config({}, num_warps=num_warps, num_stages=num_stages)
492
+ for num_warps in [2, 4, 8]
493
+ for num_stages in [2, 3, 4]
494
+ ],
495
+ key=["H", "K", "V", "BT", "BK", "BV", "IS_VARLEN"],
496
+)
497
+@triton.jit(do_not_specialize=["T"])
498
+def recompute_w_u_fwd_kernel(
499
+ q,
500
+ k,
501
+ qg,
502
+ kg,
503
+ v,
504
+ beta,
505
+ w,
506
+ u,
507
+ A,
508
+ gk,
509
+ cu_seqlens,
510
+ chunk_indices,
511
+ T,
512
+ H: tl.constexpr,
513
+ K: tl.constexpr,
514
+ V: tl.constexpr,
515
+ BT: tl.constexpr,
516
+ BK: tl.constexpr,
517
+ BV: tl.constexpr,
518
+ STORE_QG: tl.constexpr,
519
+ STORE_KG: tl.constexpr,
520
+ IS_VARLEN: tl.constexpr,
521
+ DOT_PRECISION: tl.constexpr,
522
+):
523
+ i_t, i_bh = tl.program_id(0), tl.program_id(1)
524
+ i_b, i_h = i_bh // H, i_bh % H
525
+ if IS_VARLEN:
526
+ i_n, i_t = (
527
+ tl.load(chunk_indices + i_t * 2).to(tl.int32),
528
+ tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
529
+ )
530
+ bos, eos = (
531
+ tl.load(cu_seqlens + i_n).to(tl.int32),
532
+ tl.load(cu_seqlens + i_n + 1).to(tl.int32),
533
+ )
534
+ T = eos - bos
535
+ else:
536
+ bos, eos = i_b * T, i_b * T + T
537
+ p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,))
538
+ b_b = tl.load(p_b, boundary_check=(0,))
539
+
540
+ p_A = tl.make_block_ptr(
541
+ A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)
542
+ )
543
+ b_A = tl.load(p_A, boundary_check=(0, 1))
544
+
545
+ for i_v in range(tl.cdiv(V, BV)):
546
+ p_v = tl.make_block_ptr(
547
+ v + (bos * H + i_h) * V,
548
+ (T, V),
549
+ (H * V, 1),
550
+ (i_t * BT, i_v * BV),
551
+ (BT, BV),
552
+ (1, 0),
553
+ )
554
+ p_u = tl.make_block_ptr(
555
+ u + (bos * H + i_h) * V,
556
+ (T, V),
557
+ (H * V, 1),
558
+ (i_t * BT, i_v * BV),
559
+ (BT, BV),
560
+ (1, 0),
561
+ )
562
+ b_v = tl.load(p_v, boundary_check=(0, 1))
563
+ b_vb = (b_v * b_b[:, None]).to(b_v.dtype)
564
+ b_u = tl.dot(b_A, b_vb, input_precision=DOT_PRECISION)
565
+ tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1))
566
+
567
+ for i_k in range(tl.cdiv(K, BK)):
568
+ p_w = tl.make_block_ptr(
569
+ w + (bos * H + i_h) * K,
570
+ (T, K),
571
+ (H * K, 1),
572
+ (i_t * BT, i_k * BK),
573
+ (BT, BK),
574
+ (1, 0),
575
+ )
576
+ p_k = tl.make_block_ptr(
577
+ k + (bos * H + i_h) * K,
578
+ (T, K),
579
+ (H * K, 1),
580
+ (i_t * BT, i_k * BK),
581
+ (BT, BK),
582
+ (1, 0),
583
+ )
584
+ b_k = tl.load(p_k, boundary_check=(0, 1))
585
+ b_kb = b_k * b_b[:, None]
586
+
587
+ p_gk = tl.make_block_ptr(
588
+ gk + (bos * H + i_h) * K,
589
+ (T, K),
590
+ (H * K, 1),
591
+ (i_t * BT, i_k * BK),
592
+ (BT, BK),
593
+ (1, 0),
594
+ )
595
+ b_gk = tl.load(p_gk, boundary_check=(0, 1))
596
+ b_kb *= exp2(b_gk)
597
+ if STORE_QG:
598
+ p_q = tl.make_block_ptr(
599
+ q + (bos * H + i_h) * K,
600
+ (T, K),
601
+ (H * K, 1),
602
+ (i_t * BT, i_k * BK),
603
+ (BT, BK),
604
+ (1, 0),
605
+ )
606
+ p_qg = tl.make_block_ptr(
607
+ qg + (bos * H + i_h) * K,
608
+ (T, K),
609
+ (H * K, 1),
610
+ (i_t * BT, i_k * BK),
611
+ (BT, BK),
612
+ (1, 0),
613
+ )
614
+ b_q = tl.load(p_q, boundary_check=(0, 1))
615
+ b_qg = b_q * exp2(b_gk)
616
+ tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1))
617
+ if STORE_KG:
618
+ last_idx = min(i_t * BT + BT, T) - 1
619
+
620
+ o_k = i_k * BK + tl.arange(0, BK)
621
+ m_k = o_k < K
622
+ b_gn = tl.load(
623
+ gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.0
624
+ )
625
+ b_kg = b_k * exp2(b_gn - b_gk)
626
+
627
+ p_kg = tl.make_block_ptr(
628
+ kg + (bos * H + i_h) * K,
629
+ (T, K),
630
+ (H * K, 1),
631
+ (i_t * BT, i_k * BK),
632
+ (BT, BK),
633
+ (1, 0),
634
+ )
635
+ tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1))
636
+
637
+ b_w = tl.dot(b_A, b_kb.to(b_k.dtype))
638
+ tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1))
639
+
640
+
641
+def recompute_w_u_fwd(
642
+ k: torch.Tensor,
643
+ v: torch.Tensor,
644
+ beta: torch.Tensor,
645
+ A: torch.Tensor,
646
+ q: torch.Tensor | None = None,
647
+ gk: torch.Tensor | None = None,
648
+ cu_seqlens: torch.Tensor | None = None,
649
+ chunk_indices: torch.Tensor | None = None,
650
+) -> tuple[torch.Tensor, torch.Tensor]:
651
+ B, T, H, K, V = *k.shape, v.shape[-1]
652
+ BT = A.shape[-1]
653
+ BK = 64
654
+ BV = 64
655
+
656
+ if chunk_indices is None and cu_seqlens is not None:
657
+ chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
658
+ NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
659
+
660
+ w = torch.empty_like(k)
661
+ u = torch.empty_like(v)
662
+ kg = torch.empty_like(k) if gk is not None else None
663
+ recompute_w_u_fwd_kernel[(NT, B * H)](
664
+ q=q,
665
+ k=k,
666
+ qg=None,
667
+ kg=kg,
668
+ v=v,
669
+ beta=beta,
670
+ w=w,
671
+ u=u,
672
+ A=A,
673
+ gk=gk,
674
+ cu_seqlens=cu_seqlens,
675
+ chunk_indices=chunk_indices,
676
+ T=T,
677
+ H=H,
678
+ K=K,
679
+ V=V,
680
+ BT=BT,
681
+ BK=BK,
682
+ BV=BV,
683
+ DOT_PRECISION="ieee",
684
+ )
685
+ return w, u, None, kg
686
+
687
+
688
+@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
689
+@triton.autotune(
690
+ configs=[
691
+ triton.Config({"BK": BK, "BV": BV}, num_warps=num_warps, num_stages=num_stages)
692
+ for BK in [32, 64]
693
+ for BV in [64, 128]
694
+ for num_warps in [2, 4, 8]
695
+ for num_stages in [2, 3, 4]
696
+ ],
697
+ key=["BT"],
698
+)
699
+@triton.jit(do_not_specialize=["T"])
700
+def chunk_gla_fwd_kernel_o(
701
+ q,
702
+ v,
703
+ g,
704
+ h,
705
+ o,
706
+ A,
707
+ cu_seqlens,
708
+ chunk_indices,
709
+ scale,
710
+ T,
711
+ H: tl.constexpr,
712
+ K: tl.constexpr,
713
+ V: tl.constexpr,
714
+ BT: tl.constexpr,
715
+ BK: tl.constexpr,
716
+ BV: tl.constexpr,
717
+ IS_VARLEN: tl.constexpr,
718
+):
719
+ i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
720
+ i_b, i_h = i_bh // H, i_bh % H
721
+ if IS_VARLEN:
722
+ i_tg = i_t
723
+ i_n, i_t = (
724
+ tl.load(chunk_indices + i_t * 2).to(tl.int32),
725
+ tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
726
+ )
727
+ bos, eos = (
728
+ tl.load(cu_seqlens + i_n).to(tl.int32),
729
+ tl.load(cu_seqlens + i_n + 1).to(tl.int32),
730
+ )
731
+ T = eos - bos
732
+ NT = tl.cdiv(T, BT)
733
+ else:
734
+ NT = tl.cdiv(T, BT)
735
+ i_tg = i_b * NT + i_t
736
+ bos, eos = i_b * T, i_b * T + T
737
+
738
+ m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :]
739
+
740
+ b_o = tl.zeros([BT, BV], dtype=tl.float32)
741
+ for i_k in range(tl.cdiv(K, BK)):
742
+ p_q = tl.make_block_ptr(
743
+ q + (bos * H + i_h) * K,
744
+ (T, K),
745
+ (H * K, 1),
746
+ (i_t * BT, i_k * BK),
747
+ (BT, BK),
748
+ (1, 0),
749
+ )
750
+ p_g = tl.make_block_ptr(
751
+ g + (bos * H + i_h) * K,
752
+ (T, K),
753
+ (H * K, 1),
754
+ (i_t * BT, i_k * BK),
755
+ (BT, BK),
756
+ (1, 0),
757
+ )
758
+ p_h = tl.make_block_ptr(
759
+ h + (i_tg * H + i_h) * K * V,
760
+ (V, K),
761
+ (K, 1),
762
+ (i_v * BV, i_k * BK),
763
+ (BV, BK),
764
+ (1, 0),
765
+ )
766
+
767
+ # [BT, BK]
768
+ b_q = tl.load(p_q, boundary_check=(0, 1))
769
+ b_q = (b_q * scale).to(b_q.dtype)
770
+ # [BT, BK]
771
+ b_g = tl.load(p_g, boundary_check=(0, 1))
772
+ # [BT, BK]
773
+ b_qg = (b_q * exp2(b_g)).to(b_q.dtype)
774
+ # [BV, BK]
775
+ b_h = tl.load(p_h, boundary_check=(0, 1))
776
+ # [BT, BV]
777
+ if i_k >= 0:
778
+ b_o += tl.dot(b_qg, tl.trans(b_h).to(b_qg.dtype))
779
+ p_v = tl.make_block_ptr(
780
+ v + (bos * H + i_h) * V,
781
+ (T, V),
782
+ (H * V, 1),
783
+ (i_t * BT, i_v * BV),
784
+ (BT, BV),
785
+ (1, 0),
786
+ )
787
+ p_o = tl.make_block_ptr(
788
+ o + (bos * H + i_h) * V,
789
+ (T, V),
790
+ (H * V, 1),
791
+ (i_t * BT, i_v * BV),
792
+ (BT, BV),
793
+ (1, 0),
794
+ )
795
+ p_A = tl.make_block_ptr(
796
+ A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)
797
+ )
798
+ # [BT, BV]
799
+ b_v = tl.load(p_v, boundary_check=(0, 1))
800
+ # [BT, BT]
801
+ b_A = tl.load(p_A, boundary_check=(0, 1))
802
+ b_A = tl.where(m_s, b_A, 0.0).to(b_v.dtype)
803
+ b_o += tl.dot(b_A, b_v, allow_tf32=False)
804
+ tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1))
805
+
806
+
807
+def chunk_gla_fwd_o_gk(
808
+ q: torch.Tensor,
809
+ v: torch.Tensor,
810
+ g: torch.Tensor,
811
+ A: torch.Tensor,
812
+ h: torch.Tensor,
813
+ o: torch.Tensor,
814
+ scale: float,
815
+ cu_seqlens: torch.Tensor | None = None,
816
+ chunk_indices: torch.Tensor | None = None,
817
+ chunk_size: int = FLA_CHUNK_SIZE,
818
+):
819
+ B, T, H, K, V = *q.shape, v.shape[-1]
820
+ BT = chunk_size
821
+
822
+ if chunk_indices is None and cu_seqlens is not None:
823
+ chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size)
824
+ NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
825
+
826
+ def grid(meta):
827
+ return (cdiv(V, meta["BV"]), NT, B * H)
828
+
829
+ chunk_gla_fwd_kernel_o[grid](
830
+ q=q,
831
+ v=v,
832
+ g=g,
833
+ h=h,
834
+ o=o,
835
+ A=A,
836
+ cu_seqlens=cu_seqlens,
837
+ chunk_indices=chunk_indices,
838
+ scale=scale,
839
+ T=T,
840
+ H=H,
841
+ K=K,
842
+ V=V,
843
+ BT=BT,
844
+ )
845
+ return o
846
+
847
+
848
+@triton.heuristics(
849
+ {
850
+ "HAS_BIAS": lambda args: args["g_bias"] is not None,
851
+ "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
852
+ }
853
+)
854
+@triton.autotune(
855
+ configs=[
856
+ triton.Config({"BD": BD}, num_warps=num_warps)
857
+ for BD in [32, 64]
858
+ for num_warps in [2, 4, 8]
859
+ ],
860
+ key=["H", "D", "BT", "IS_VARLEN"],
861
+)
862
+@triton.jit(do_not_specialize=["T"])
863
+def kda_gate_cumsum_fwd_kernel(
864
+ g,
865
+ A,
866
+ y,
867
+ g_bias,
868
+ cu_seqlens,
869
+ chunk_indices,
870
+ cumsum_scale,
871
+ beta,
872
+ threshold,
873
+ SAFE_GATE: tl.constexpr,
874
+ LOWER_BOUND: tl.constexpr,
875
+ T,
876
+ H: tl.constexpr,
877
+ D: tl.constexpr,
878
+ BT: tl.constexpr,
879
+ BD: tl.constexpr,
880
+ HAS_BIAS: tl.constexpr,
881
+ IS_VARLEN: tl.constexpr,
882
+):
883
+ i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
884
+ i_b, i_h = i_bh // H, i_bh % H
885
+ if IS_VARLEN:
886
+ i_n, i_t = (
887
+ tl.load(chunk_indices + i_t * 2).to(tl.int32),
888
+ tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
889
+ )
890
+ bos, eos = (
891
+ tl.load(cu_seqlens + i_n).to(tl.int32),
892
+ tl.load(cu_seqlens + i_n + 1).to(tl.int32),
893
+ )
894
+ T = eos - bos
895
+ else:
896
+ bos = i_b * T
897
+
898
+ p_g = tl.make_block_ptr(
899
+ g + (bos * H + i_h) * D,
900
+ (T, D),
901
+ (H * D, 1),
902
+ (i_t * BT, i_d * BD),
903
+ (BT, BD),
904
+ (1, 0),
905
+ )
906
+ p_y = tl.make_block_ptr(
907
+ y + (bos * H + i_h) * D,
908
+ (T, D),
909
+ (H * D, 1),
910
+ (i_t * BT, i_d * BD),
911
+ (BT, BD),
912
+ (1, 0),
913
+ )
914
+
915
+ b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32)
916
+ if HAS_BIAS:
917
+ o_d = i_d * BD + tl.arange(0, BD)
918
+ b_bias = tl.load(g_bias + i_h * D + o_d, mask=o_d < D, other=0.0).to(tl.float32)
919
+ b_g = b_g + b_bias[None, :]
920
+
921
+ b_a = tl.load(A + i_h).to(tl.float32)
922
+ b_a = tl.exp(b_a) if SAFE_GATE else -tl.exp(b_a)
923
+ if SAFE_GATE:
924
+ # y = lower_bound * sigmoid(exp(A) * (g + g_bias)), bounded to
925
+ # (lower_bound, 0) for safe-gate checkpoints.
926
+ b_gate = LOWER_BOUND / (1.0 + tl.exp(-(b_a * b_g)))
927
+ else:
928
+ b_g_scaled = b_g * beta
929
+ b_softplus = tl.where(
930
+ b_g_scaled > threshold,
931
+ b_g,
932
+ (1.0 / beta) * log(1.0 + tl.exp(b_g_scaled)),
933
+ )
934
+ b_gate = b_a * b_softplus
935
+
936
+ # Out-of-bounds rows (load returns 0, but softplus/bias can still make
937
+ # b_gate non-zero) participate in the dot product. They only contribute to
938
+ # out-of-bounds output rows, which are masked away by `boundary_check` on
939
+ # the store, so visible output matches unfused gate + chunk-local cumsum.
940
+ o_t = tl.arange(0, BT)
941
+ m_cumsum = tl.where(o_t[:, None] >= o_t[None, :], 1.0, 0.0)
942
+ b_y = tl.dot(m_cumsum, b_gate, allow_tf32=False) * cumsum_scale
943
+ tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1))
944
+
945
+
946
+def fused_kda_gate_chunk_cumsum(
947
+ raw_g: torch.Tensor,
948
+ A_log: torch.Tensor,
949
+ g_bias: torch.Tensor | None = None,
950
+ beta: float = 1.0,
951
+ threshold: float = 20.0,
952
+ cu_seqlens: torch.Tensor | None = None,
953
+ chunk_indices: torch.Tensor | None = None,
954
+ chunk_size: int = FLA_CHUNK_SIZE,
955
+ output_dtype: torch.dtype | None = torch.float,
956
+ safe_gate: bool = False,
957
+ lower_bound: float = -5.0,
958
+) -> torch.Tensor:
959
+ if cu_seqlens is not None:
960
+ assert raw_g.shape[0] == 1, (
961
+ "Only batch size 1 is supported when cu_seqlens are provided"
962
+ )
963
+ B, T, H, D = raw_g.shape
964
+ if chunk_indices is None and cu_seqlens is not None:
965
+ chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size)
966
+ NT = cdiv(T, chunk_size) if cu_seqlens is None else len(chunk_indices)
967
+
968
+ A_log = A_log.reshape(-1)
969
+ if g_bias is not None:
970
+ g_bias = g_bias.reshape(-1)
971
+ y = torch.empty_like(raw_g, dtype=output_dtype or raw_g.dtype)
972
+
973
+ def grid(meta):
974
+ return (cdiv(meta["D"], meta["BD"]), NT, B * H)
975
+
976
+ kda_gate_cumsum_fwd_kernel[grid](
977
+ g=raw_g,
978
+ A=A_log,
979
+ y=y,
980
+ g_bias=g_bias,
981
+ cu_seqlens=cu_seqlens,
982
+ chunk_indices=chunk_indices,
983
+ # RCP_LN2 folds in the natural-log -> log2 conversion so downstream
984
+ # exp2-based kernels reproduce exp(g). Keep this in sync with the
985
+ # `use_exp2=True` path in `_chunk_kda_fwd_with_cumulative_g`.
986
+ cumsum_scale=RCP_LN2,
987
+ beta=beta,
988
+ threshold=threshold,
989
+ SAFE_GATE=safe_gate,
990
+ LOWER_BOUND=lower_bound,
991
+ T=T,
992
+ H=H,
993
+ D=D,
994
+ BT=chunk_size,
995
+ )
996
+ return y
997
+
998
+
999
+def _chunk_kda_fwd_with_cumulative_g(
1000
+ q: torch.Tensor,
1001
+ k: torch.Tensor,
1002
+ v: torch.Tensor,
1003
+ g: torch.Tensor,
1004
+ beta: torch.Tensor,
1005
+ scale: float,
1006
+ initial_state: torch.Tensor,
1007
+ output_final_state: bool,
1008
+ cu_seqlens: torch.Tensor | None = None,
1009
+ chunk_indices: torch.Tensor | None = None,
1010
+ chunk_size: int = FLA_CHUNK_SIZE,
1011
+):
1012
+ # `g` must already be chunk-local cumulatively-summed AND scaled by
1013
+ # RCP_LN2 (so the downstream exp2-based kernels reproduce exp(g)).
1014
+ # Use `chunk_kda_fwd` or `chunk_kda_with_fused_gate_fwd` instead of
1015
+ # calling this helper directly unless that invariant is upheld.
1016
+ # the intra Aqk is kept in fp32
1017
+ # the computation has very marginal effect on the entire throughput
1018
+ A, Aqk = chunk_kda_scaled_dot_kkt_fwd(
1019
+ q=q,
1020
+ k=k,
1021
+ gk=g,
1022
+ beta=beta,
1023
+ scale=scale,
1024
+ cu_seqlens=cu_seqlens,
1025
+ chunk_indices=chunk_indices,
1026
+ chunk_size=chunk_size,
1027
+ output_dtype=torch.float32,
1028
+ )
1029
+ A = solve_tril(A=A, cu_seqlens=cu_seqlens, output_dtype=k.dtype)
1030
+ w, u, _, kg = recompute_w_u_fwd(
1031
+ k=k,
1032
+ v=v,
1033
+ beta=beta,
1034
+ A=A,
1035
+ gk=g,
1036
+ cu_seqlens=cu_seqlens,
1037
+ chunk_indices=chunk_indices,
1038
+ )
1039
+ del A
1040
+ h, v_new, final_state = chunk_gated_delta_rule_fwd_h(
1041
+ k=kg,
1042
+ w=w,
1043
+ u=u,
1044
+ gk=g,
1045
+ initial_state=initial_state,
1046
+ output_final_state=output_final_state,
1047
+ cu_seqlens=cu_seqlens,
1048
+ chunk_indices=chunk_indices,
1049
+ use_exp2=True,
1050
+ )
1051
+ del w, u, kg
1052
+ o = chunk_gla_fwd_o_gk(
1053
+ q=q,
1054
+ v=v_new,
1055
+ g=g,
1056
+ A=Aqk,
1057
+ h=h,
1058
+ o=v,
1059
+ scale=scale,
1060
+ cu_seqlens=cu_seqlens,
1061
+ chunk_indices=chunk_indices,
1062
+ chunk_size=chunk_size,
1063
+ )
1064
+ del Aqk, v_new, h
1065
+ return o, final_state
1066
+
1067
+
1068
+def chunk_kda_fwd(
1069
+ q: torch.Tensor,
1070
+ k: torch.Tensor,
1071
+ v: torch.Tensor,
1072
+ g: torch.Tensor,
1073
+ beta: torch.Tensor,
1074
+ scale: float,
1075
+ initial_state: torch.Tensor,
1076
+ output_final_state: bool,
1077
+ cu_seqlens: torch.Tensor | None = None,
1078
+):
1079
+ chunk_size = FLA_CHUNK_SIZE
1080
+ chunk_indices = (
1081
+ prepare_chunk_indices(cu_seqlens, chunk_size)
1082
+ if cu_seqlens is not None
1083
+ else None
1084
+ )
1085
+ g = chunk_local_cumsum(
1086
+ g,
1087
+ chunk_size=chunk_size,
1088
+ cu_seqlens=cu_seqlens,
1089
+ chunk_indices=chunk_indices,
1090
+ )
1091
+ # KDA evaluates cumulative gate decays with exp2. Convert from natural-log
1092
+ # space so exp(x) is preserved as exp2(x / ln(2)).
1093
+ g = g * RCP_LN2
1094
+ return _chunk_kda_fwd_with_cumulative_g(
1095
+ q=q,
1096
+ k=k,
1097
+ v=v,
1098
+ g=g,
1099
+ beta=beta,
1100
+ scale=scale,
1101
+ initial_state=initial_state,
1102
+ output_final_state=output_final_state,
1103
+ cu_seqlens=cu_seqlens,
1104
+ chunk_indices=chunk_indices,
1105
+ chunk_size=chunk_size,
1106
+ )
1107
+
1108
+
1109
+def chunk_kda_with_fused_gate_fwd(
1110
+ q: torch.Tensor,
1111
+ k: torch.Tensor,
1112
+ v: torch.Tensor,
1113
+ raw_g: torch.Tensor,
1114
+ beta: torch.Tensor,
1115
+ A_log: torch.Tensor,
1116
+ g_bias: torch.Tensor | None,
1117
+ scale: float,
1118
+ initial_state: torch.Tensor,
1119
+ output_final_state: bool,
1120
+ cu_seqlens: torch.Tensor | None = None,
1121
+ safe_gate: bool = False,
1122
+ lower_bound: float = -5.0,
1123
+):
1124
+ chunk_size = FLA_CHUNK_SIZE
1125
+ chunk_indices = (
1126
+ prepare_chunk_indices(cu_seqlens, chunk_size)
1127
+ if cu_seqlens is not None
1128
+ else None
1129
+ )
1130
+ g = fused_kda_gate_chunk_cumsum(
1131
+ raw_g,
1132
+ A_log=A_log,
1133
+ g_bias=g_bias,
1134
+ cu_seqlens=cu_seqlens,
1135
+ chunk_indices=chunk_indices,
1136
+ chunk_size=chunk_size,
1137
+ safe_gate=safe_gate,
1138
+ lower_bound=lower_bound,
1139
+ )
1140
+ return _chunk_kda_fwd_with_cumulative_g(
1141
+ q=q,
1142
+ k=k,
1143
+ v=v,
1144
+ g=g,
1145
+ beta=beta,
1146
+ scale=scale,
1147
+ initial_state=initial_state,
1148
+ output_final_state=output_final_state,
1149
+ cu_seqlens=cu_seqlens,
1150
+ chunk_indices=chunk_indices,
1151
+ chunk_size=chunk_size,
1152
+ )
1153
+
1154
+
1155
+def chunk_kda(
1156
+ q: torch.Tensor,
1157
+ k: torch.Tensor,
1158
+ v: torch.Tensor,
1159
+ g: torch.Tensor,
1160
+ beta: torch.Tensor,
1161
+ scale: float = None,
1162
+ initial_state: torch.Tensor = None,
1163
+ output_final_state: bool = False,
1164
+ use_qk_l2norm_in_kernel: bool = False,
1165
+ cu_seqlens: torch.Tensor | None = None,
1166
+ **kwargs,
1167
+):
1168
+ if scale is None:
1169
+ scale = k.shape[-1] ** -0.5
1170
+
1171
+ if use_qk_l2norm_in_kernel:
1172
+ q = l2norm_fwd(q.contiguous())
1173
+ k = l2norm_fwd(k.contiguous())
1174
+
1175
+ o, final_state = chunk_kda_fwd(
1176
+ q=q,
1177
+ k=k,
1178
+ v=v.contiguous(),
1179
+ g=g.contiguous(),
1180
+ beta=beta.contiguous(),
1181
+ scale=scale,
1182
+ initial_state=initial_state.contiguous(),
1183
+ output_final_state=output_final_state,
1184
+ cu_seqlens=cu_seqlens,
1185
+ )
1186
+ return o, final_state
1187
+
1188
+
1189
+def chunk_kda_with_fused_gate(
1190
+ q: torch.Tensor,
1191
+ k: torch.Tensor,
1192
+ v: torch.Tensor,
1193
+ raw_g: torch.Tensor,
1194
+ beta: torch.Tensor,
1195
+ A_log: torch.Tensor,
1196
+ g_bias: torch.Tensor | None,
1197
+ scale: float | None = None,
1198
+ initial_state: torch.Tensor | None = None,
1199
+ output_final_state: bool = False,
1200
+ use_qk_l2norm_in_kernel: bool = False,
1201
+ cu_seqlens: torch.Tensor | None = None,
1202
+ safe_gate: bool = False,
1203
+ lower_bound: float = -5.0,
1204
+ **kwargs,
1205
+):
1206
+ """Run chunk KDA from raw gate projection using fused gate+cumsum."""
1207
+ if scale is None:
1208
+ scale = k.shape[-1] ** -0.5
1209
+
1210
+ if use_qk_l2norm_in_kernel:
1211
+ q = l2norm_fwd(q.contiguous())
1212
+ k = l2norm_fwd(k.contiguous())
1213
+
1214
+ o, final_state = chunk_kda_with_fused_gate_fwd(
1215
+ q=q,
1216
+ k=k,
1217
+ v=v.contiguous(),
1218
+ raw_g=raw_g.contiguous(),
1219
+ beta=beta.contiguous(),
1220
+ A_log=A_log,
1221
+ g_bias=g_bias,
1222
+ scale=scale,
1223
+ initial_state=initial_state.contiguous() if initial_state is not None else None,
1224
+ output_final_state=output_final_state,
1225
+ cu_seqlens=cu_seqlens,
1226
+ safe_gate=safe_gate,
1227
+ lower_bound=lower_bound,
1228
+ )
1229
+ return o, final_state
1230
+
1231
+
1232
+@triton.autotune(
1233
+ configs=[
1234
+ triton.Config({"BT": bt}, num_warps=nw, num_stages=ns)
1235
+ for bt in BT_LIST_AUTOTUNE
1236
+ for nw in NUM_WARPS_AUTOTUNE
1237
+ for ns in [2, 3]
1238
+ ],
1239
+ key=["H", "D"],
1240
+)
1241
+@triton.jit
1242
+def kda_gate_fwd_kernel(
1243
+ g,
1244
+ A,
1245
+ y,
1246
+ g_bias,
1247
+ beta: tl.constexpr,
1248
+ threshold: tl.constexpr,
1249
+ SAFE_GATE: tl.constexpr,
1250
+ LOWER_BOUND: tl.constexpr,
1251
+ T,
1252
+ H,
1253
+ D: tl.constexpr,
1254
+ BT: tl.constexpr,
1255
+ BD: tl.constexpr,
1256
+ HAS_BIAS: tl.constexpr,
1257
+):
1258
+ i_t, i_h = tl.program_id(0), tl.program_id(1)
1259
+ n_t = i_t * BT
1260
+
1261
+ b_a = tl.load(A + i_h).to(tl.float32)
1262
+ b_a = tl.exp(b_a) if SAFE_GATE else -tl.exp(b_a)
1263
+
1264
+ stride_row = H * D
1265
+ stride_col = 1
1266
+
1267
+ g_ptr = tl.make_block_ptr(
1268
+ base=g + i_h * D,
1269
+ shape=(T, D),
1270
+ strides=(stride_row, stride_col),
1271
+ offsets=(n_t, 0),
1272
+ block_shape=(BT, BD),
1273
+ order=(1, 0),
1274
+ )
1275
+
1276
+ y_ptr = tl.make_block_ptr(
1277
+ base=y + i_h * D,
1278
+ shape=(T, D),
1279
+ strides=(stride_row, stride_col),
1280
+ offsets=(n_t, 0),
1281
+ block_shape=(BT, BD),
1282
+ order=(1, 0),
1283
+ )
1284
+
1285
+ b_g = tl.load(g_ptr, boundary_check=(0, 1)).to(tl.float32)
1286
+
1287
+ if HAS_BIAS:
1288
+ n_d = tl.arange(0, BD)
1289
+ bias_mask = n_d < D
1290
+ b_bias = tl.load(g_bias + i_h * D + n_d, mask=bias_mask, other=0.0).to(
1291
+ tl.float32
1292
+ )
1293
+ b_g = b_g + b_bias[None, :]
1294
+
1295
+ if SAFE_GATE:
1296
+ # y = lower_bound * sigmoid(exp(A) * (g + g_bias)), bounded to
1297
+ # (lower_bound, 0) for safe-gate checkpoints.
1298
+ b_y = LOWER_BOUND / (1.0 + tl.exp(-(b_a * b_g)))
1299
+ else:
1300
+ # softplus(x, beta) = (1/beta) * log(1 + exp(beta * x))
1301
+ # When beta * x > threshold, use linear approximation x
1302
+ # Use threshold to switch to linear when beta*x > threshold
1303
+ g_scaled = b_g * beta
1304
+ use_linear = g_scaled > threshold
1305
+ sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled)))
1306
+ b_y = b_a * sp
1307
+
1308
+ tl.store(y_ptr, b_y.to(y.dtype.element_ty), boundary_check=(0, 1))
1309
+
1310
+
1311
+def fused_kda_gate(
1312
+ g: torch.Tensor,
1313
+ A: torch.Tensor,
1314
+ head_k_dim: int,
1315
+ g_bias: torch.Tensor | None = None,
1316
+ beta: float = 1.0,
1317
+ threshold: float = 20.0,
1318
+ safe_gate: bool = False,
1319
+ lower_bound: float | None = -5.0,
1320
+) -> torch.Tensor:
1321
+ """
1322
+ Forward pass for KDA gate:
1323
+ input g: [..., H*D]
1324
+ param A: [H] or [1, 1, H, 1]
1325
+ beta: softplus beta parameter (softplus branch only)
1326
+ threshold: softplus threshold parameter (softplus branch only)
1327
+ safe_gate: when False (default) compute y = -exp(A)*softplus(g+g_bias);
1328
+ when True compute the bounded y = lower_bound*sigmoid(exp(A)*(g+g_bias))
1329
+ lower_bound: floor for the safe_gate branch (default -5.0)
1330
+ return : [..., H, D]
1331
+ """
1332
+ orig_shape = g.shape[:-1]
1333
+
1334
+ g = g.view(-1, g.shape[-1])
1335
+ T = g.shape[0]
1336
+ HD = g.shape[1]
1337
+ H = A.numel()
1338
+ assert H * head_k_dim == HD
1339
+
1340
+ y = torch.empty_like(g, dtype=torch.float32)
1341
+
1342
+ def grid(meta):
1343
+ return (cdiv(T, meta["BT"]), H)
1344
+
1345
+ kda_gate_fwd_kernel[grid](
1346
+ g,
1347
+ A,
1348
+ y,
1349
+ g_bias,
1350
+ beta,
1351
+ threshold,
1352
+ safe_gate,
1353
+ lower_bound if lower_bound is not None else -5.0,
1354
+ T,
1355
+ H,
1356
+ head_k_dim,
1357
+ BD=next_power_of_2(head_k_dim),
1358
+ HAS_BIAS=g_bias is not None,
1359
+ )
1360
+
1361
+ y = y.view(*orig_shape, H, head_k_dim)
1362
+ return y
vllm/models/glm5next/nvidia/attention.py
new
+590
@@ -0,0 +1,590 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+import torch.nn.functional as F
6
+from torch import nn
7
+
8
+from vllm.config import (
9
+ CacheConfig,
10
+ VllmConfig,
11
+)
12
+from vllm.distributed import (
13
+ get_tensor_model_parallel_world_size,
14
+)
15
+from vllm.logger import init_logger
16
+from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm
17
+from vllm.model_executor.layers.linear import (
18
+ ColumnParallelLinear,
19
+ MergedColumnParallelLinear,
20
+ ReplicatedLinear,
21
+ RowParallelLinear,
22
+)
23
+from vllm.model_executor.layers.mla import MLAModules, MultiHeadLatentAttentionWrapper
24
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
25
+from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding, get_rope
26
+from vllm.model_executor.layers.sparse_attn_indexer_kpool import SparseAttnIndexerKpool
27
+from vllm.model_executor.models.deepseek_v2 import (
28
+ DeepSeekV2FusedQkvAProjLinear,
29
+ DeepseekV32IndexerCache,
30
+ yarn_get_mscale,
31
+)
32
+from vllm.model_executor.utils import maybe_disable_graph_partition
33
+from vllm.models.glm5next.nvidia.ops.kpool_compress import fwht128_quant_fp8
34
+from vllm.platforms import current_platform
35
+from vllm.transformers_utils.configs.glm5_next import Glm5NextConfig
36
+from vllm.utils.deep_gemm import PAGED_MQA_PAGE_SIZES
37
+from vllm.v1.kv_cache_interface import KpoolTailSpec, MLAAttentionSpec
38
+
39
+logger = init_logger(__name__)
40
+
41
+# Shared torch.compile config for the indexer's small-kernel leaves. The MLA
42
+# indexer runs under breakable-CG (CompilationMode.NONE), which blocks FX-graph
43
+# fusion of the surrounding eager ops; carving each cluster into its own
44
+# @torch.compile leaf (backend==inductor) still fuses them. Matches the
45
+# grouped_topk / _cast_sigmoid leaf pattern.
46
+_INDEXER_COMPILE = dict(
47
+ dynamic=True,
48
+ backend=current_platform.simple_compile_backend,
49
+ options=maybe_disable_graph_partition(current_platform.simple_compile_backend),
50
+)
51
+
52
+
53
+@torch.compile(**_INDEXER_COMPILE)
54
+def _fused_indexer_k_norm(
55
+ x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, dim: int, eps: float
56
+) -> torch.Tensor:
57
+ # Fuse fp32 cast + layer_norm + cast-back (was 3 kernels) into one.
58
+ return F.layer_norm(x.float(), (dim,), weight, bias, eps).type_as(x)
59
+
60
+
61
+@torch.compile(**_INDEXER_COMPILE)
62
+def _fused_indexer_weight_scale(
63
+ weights: torch.Tensor, q_scale: torch.Tensor, scale: float
64
+) -> torch.Tensor:
65
+ # Fuse the weight-scaling muls (was 2 kernels) into one. `scale` folds
66
+ # softmax_scale (head_dim**-0.5) and n_head**-0.5 into a single constant.
67
+ return (weights.unsqueeze(-1) * q_scale * scale).squeeze(-1)
68
+
69
+
70
+@torch.compile(**_INDEXER_COMPILE)
71
+def _pad_indexer_heads(x: torch.Tensor, pad: int) -> torch.Tensor:
72
+ # DeepGEMM MQA-logits needs num_heads in {32,64}; zero-pad the head dim.
73
+ # Fuse new_zeros + cat (was 2 kernels) into one. Pad values are zero (exact
74
+ # in fp8 e4m3 and zero-weight in the logits sum), so numerically a no-op.
75
+ return torch.cat([x, x.new_zeros(x.shape[0], pad, *x.shape[2:])], dim=1)
76
+
77
+
78
+class Glm5NextIndexerCache(DeepseekV32IndexerCache):
79
+ """Indexer K cache that stores kpool-compressed entries.
80
+
81
+ Setting ``tokens_per_state = index_kpool`` on the KV cache spec makes vLLM's
82
+ indexer metadata builder emit pool-granular ``slot_mapping`` /
83
+ ``seq_lens`` / ``cu_seq_lens`` / ``page_table`` for free, and shrinks the
84
+ cache allocation store one state per ``index_kpool`` tokens. The pool
85
+ *content* (softmax-weighted sum vs keep-every-Nth) is computed by the
86
+ kpool compress kernel inside the indexer op — the cache only provides the
87
+ addressing, which is identical for both schemes.
88
+
89
+ The indexer shares one block with the co-located MLA (a single
90
+ ``MLAAttentionSpec`` / block_table), so ``block_size`` is the model-wide
91
+ ``cache_config.block_size``. DeepGEMM's paged-MQA kernel
92
+ (``csrc/apis/attention.hpp``) requires ``block_kv`` to be exactly 32 or
93
+ 64, so the storage block is virtually split into pool pages of the
94
+ largest such size that tiles it (``storage_kernel_block_size``); this
95
+ needs ``block_size`` to be a multiple of ``index_kpool * 32`` (512 for
96
+ ``index_kpool = 16``). A smaller block (e.g. the default 64) silently
97
+ collapses ``storage_block_size`` (64 // 16 = 4) and only fails later at
98
+ the opaque C++ assert; ``get_kv_cache_spec`` guards this up front
99
+ instead.
100
+ """
101
+
102
+ def __init__(
103
+ self,
104
+ *,
105
+ head_dim: int,
106
+ dtype: torch.dtype,
107
+ prefix: str,
108
+ cache_config,
109
+ index_kpool: int,
110
+ ):
111
+ super().__init__(
112
+ head_dim=head_dim, dtype=dtype, prefix=prefix, cache_config=cache_config
113
+ )
114
+ assert index_kpool > 1, "Glm5NextIndexerCache expects index_kpool > 1"
115
+ # Keep chunked-prefill boundaries aligned to complete pools.
116
+ assert cache_config.block_size % index_kpool == 0, (
117
+ "Glm5NextIndexerCache: cache_config.block_size "
118
+ f"({cache_config.block_size}) must be a multiple of index_kpool "
119
+ f"({index_kpool}) so chunked-prefill boundaries stay pool-aligned."
120
+ )
121
+ self._index_kpool = index_kpool
122
+
123
+ def get_kv_cache_spec(self, vllm_config: VllmConfig):
124
+ from dataclasses import replace
125
+
126
+ spec = super().get_kv_cache_spec(vllm_config)
127
+ # ``tokens_per_state`` is the KV-spec representation of kpool
128
+ # compression in the current cache-layout API.
129
+ assert isinstance(spec, MLAAttentionSpec)
130
+ spec = replace(spec, tokens_per_state=self._index_kpool)
131
+
132
+ # DeepGEMM paged-MQA takes block_kv in {32, 64}; the storage block
133
+ # (= block_size // index_kpool) is virtually split into pool pages of
134
+ # the largest such size that tiles it, so it must be a multiple of 32.
135
+ storage_block_size = spec.block_size // self._index_kpool
136
+ assert (
137
+ spec.block_size % self._index_kpool == 0 and storage_block_size % 32 == 0
138
+ ), (
139
+ "Glm5NextIndexerCache: kpool indexer requires cache block_size to "
140
+ f"be a multiple of index_kpool * 32 ({self._index_kpool * 32}) so "
141
+ "that DeepGEMM paged-MQA pool pages (32 or 64 entries) tile the "
142
+ f"storage block, got block_size={spec.block_size} -> "
143
+ f"storage_block_size={storage_block_size}."
144
+ )
145
+ max_page_size = max(PAGED_MQA_PAGE_SIZES)
146
+ min_page_size = min(PAGED_MQA_PAGE_SIZES)
147
+ if storage_block_size <= max_page_size:
148
+ page_size = storage_block_size
149
+ elif storage_block_size % max_page_size == 0:
150
+ page_size = max_page_size
151
+ else:
152
+ page_size = min_page_size
153
+ return replace(
154
+ spec,
155
+ storage_block_size=page_size * self._index_kpool,
156
+ )
157
+
158
+
159
+class Glm5NextTailCache(DeepseekV32IndexerCache):
160
+ """Paged circular buffer for the kpool indexer's in-progress (tail) pool.
161
+
162
+ Holds the trailing incomplete pool's raw K + gate score: one block of
163
+ ``index_kpool`` slots per request, overwritten in place by ``pos % kpool``
164
+ as decode/spec-decode advances. Prefill seeds it (instead of discarding the
165
+ tail raw K+gate); the connector transfers it across PD; decode reads it to
166
+ compress the boundary pool correctly. ``KpoolTailSpec`` /
167
+ ``KpoolTailManager`` provide the no-prune, 1-block/req allocation that lets
168
+ the in-progress pool survive across steps and across transfer.
169
+
170
+ Stores raw bf16 K (``head_dim``) as the "K" half of each block and the
171
+ bf16 gate score (``head_dim``) as the "V" half -- not the fp8-compressed
172
+ entry, which lives in ``Glm5NextIndexerCache``.
173
+ """
174
+
175
+ def __init__(
176
+ self,
177
+ *,
178
+ head_dim: int,
179
+ dtype: torch.dtype,
180
+ prefix: str,
181
+ cache_config,
182
+ index_kpool: int,
183
+ ):
184
+ super().__init__(
185
+ head_dim=head_dim, dtype=dtype, prefix=prefix, cache_config=cache_config
186
+ )
187
+ assert index_kpool > 1, "Glm5NextTailCache expects index_kpool > 1"
188
+ self._index_kpool = index_kpool
189
+
190
+ def get_kv_cache_spec(self, vllm_config: VllmConfig):
191
+ # The two head slots form [K, gate score] in the generic
192
+ # [block, head, state, content] cache view.
193
+ return KpoolTailSpec(
194
+ block_size=self._index_kpool,
195
+ num_kv_heads=2,
196
+ head_size=self.head_dim,
197
+ head_size_v=0,
198
+ dtype=torch.bfloat16,
199
+ sliding_window=self._index_kpool,
200
+ )
201
+
202
+ def get_attn_backend(self):
203
+ from vllm.v1.attention.backends.mla.indexer import KpoolTailBackend
204
+
205
+ return KpoolTailBackend
206
+
207
+
208
+class Indexer(nn.Module):
209
+ def __init__(
210
+ self,
211
+ vllm_config: VllmConfig,
212
+ config: Glm5NextConfig,
213
+ hidden_size: int,
214
+ q_lora_rank: int,
215
+ quant_config: QuantizationConfig | None,
216
+ cache_config: CacheConfig | None,
217
+ topk_indices_buffer: torch.Tensor | None,
218
+ prefix: str = "",
219
+ ):
220
+ super().__init__()
221
+ self.vllm_config = vllm_config
222
+ self.config = config
223
+ self.quant_config = quant_config
224
+ # self.indexer_cfg = config.attn_module_list_cfg[0]["attn_index"]
225
+ # Indexer is only constructed for v32 configs, where these sparse-indexer
226
+ # fields are guaranteed populated; narrow away the `int | None` declared
227
+ # on Glm5NextConfig for the optional-indexer case.
228
+ assert config.index_topk is not None
229
+ assert config.index_n_heads is not None
230
+ assert config.index_head_dim is not None
231
+ assert config.index_kpool is not None
232
+ self.topk_tokens = config.index_topk
233
+ self.n_head = config.index_n_heads # 64
234
+ self.head_dim = config.index_head_dim # 128
235
+ self.rope_dim = config.qk_rope_head_dim # 64
236
+ self.index_kpool = config.index_kpool
237
+ self.q_lora_rank = q_lora_rank # 1536
238
+
239
+ # kpool
240
+ self.index_kpool_compress_ape = nn.Parameter(
241
+ torch.zeros(self.index_kpool, self.head_dim, dtype=torch.float32)
242
+ )
243
+ # Keep the checkpoint name ``index_kpool_compress_gate`` without a
244
+ # ``.weight`` suffix. F.linear consumes its [head_dim, hidden_size] shape.
245
+ self.index_kpool_compress_gate = nn.Parameter(
246
+ torch.empty(self.head_dim, hidden_size, dtype=torch.bfloat16)
247
+ )
248
+
249
+ # no tensor parallel, just replicated
250
+ self.wq_b = ReplicatedLinear(
251
+ self.q_lora_rank,
252
+ self.head_dim * self.n_head,
253
+ bias=False,
254
+ quant_config=quant_config,
255
+ prefix=f"{prefix}.wq_b",
256
+ )
257
+ # Fused wk + weights_proj: single GEMM producing [head_dim + n_head].
258
+ # FP8 wk weights are upcasted to BF16 during loading to maintain fusion.
259
+ self.wk_weights_proj = MergedColumnParallelLinear(
260
+ hidden_size,
261
+ [self.head_dim, self.n_head],
262
+ bias=False,
263
+ quant_config=None,
264
+ disable_tp=True,
265
+ prefix=f"{prefix}.wk_weights_proj",
266
+ )
267
+ self.k_norm = LayerNorm(self.head_dim, eps=1e-6)
268
+ self.softmax_scale = self.head_dim**-0.5
269
+
270
+ # Hadamard-128 rotation of the indexer query is fused with the FP8
271
+ # quant (see forward: fwht128_quant_fp8) -- no precomputed matrix.
272
+
273
+ self.scale_fmt = "ue8m0"
274
+ self.quant_block_size = 128 # TODO: get from config
275
+ self.topk_indices_buffer = topk_indices_buffer
276
+ self._wp_fp32: torch.Tensor | None = None
277
+
278
+ # NOTE: (zyongye) we use fp8 naive cache,
279
+ # where we store value in fp8 and scale in fp32
280
+ # per self.quant_block_size element
281
+ self.k_cache = Glm5NextIndexerCache(
282
+ head_dim=self.head_dim + self.head_dim // self.quant_block_size * 4,
283
+ dtype=torch.uint8,
284
+ prefix=f"{prefix}.k_cache",
285
+ cache_config=cache_config,
286
+ index_kpool=self.index_kpool,
287
+ )
288
+ # Paged tail cache (in-progress pool's raw K + gate score). Written by
289
+ # prefill (seeds the boundary pool) and decode (per-step stash); read by
290
+ # the decode kernel to compress the boundary pool. Transferred across PD
291
+ # so the decode side sees the prefill tail. See KpoolTailSpec/Manager.
292
+ self.tail_cache = Glm5NextTailCache(
293
+ head_dim=self.head_dim,
294
+ dtype=torch.bfloat16,
295
+ prefix=f"{prefix}.tail_cache",
296
+ cache_config=cache_config,
297
+ index_kpool=self.index_kpool,
298
+ )
299
+ self.max_model_len = vllm_config.model_config.max_model_len
300
+ self.prefix = prefix
301
+ from vllm.v1.attention.backends.mla.indexer import get_max_prefill_buffer_size
302
+
303
+ self.max_total_seq_len = get_max_prefill_buffer_size(vllm_config)
304
+ self.indexer_op = SparseAttnIndexerKpool(
305
+ self.k_cache,
306
+ self.quant_block_size,
307
+ self.scale_fmt,
308
+ self.topk_tokens,
309
+ self.head_dim,
310
+ self.max_model_len,
311
+ self.max_total_seq_len,
312
+ self.topk_indices_buffer,
313
+ tail_cache=self.tail_cache,
314
+ )
315
+
316
+ def forward(
317
+ self, hidden_states: torch.Tensor, qr: torch.Tensor, positions, rotary_emb
318
+ ) -> torch.Tensor:
319
+ q, _ = self.wq_b(qr)
320
+ q = q.view(-1, self.n_head, self.head_dim)
321
+
322
+ # Compute the head gate in fp32; bf16 error can change near-tie pool
323
+ # rankings on long-context tasks. Cache it after weights are loaded.
324
+ kw, _ = self.wk_weights_proj(hidden_states)
325
+ k = kw[:, : self.head_dim]
326
+ if self._wp_fp32 is None:
327
+ self._wp_fp32 = (
328
+ self.wk_weights_proj.weight.data[self.head_dim :, :]
329
+ .t()
330
+ .contiguous()
331
+ .float()
332
+ )
333
+ weights = torch.mm(hidden_states.float(), self._wp_fp32)
334
+
335
+ k = _fused_indexer_k_norm(
336
+ k, self.k_norm.weight, self.k_norm.bias, self.head_dim, self.k_norm.eps
337
+ )
338
+
339
+ if self.rope_dim > 0:
340
+ q_pe, q_nope = torch.split(
341
+ q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1
342
+ )
343
+ k_pe, k_nope = torch.split(
344
+ k, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1
345
+ )
346
+
347
+ q_pe, k_pe = rotary_emb(positions, q_pe, k_pe.unsqueeze(1))
348
+ # Note: RoPE (NeoX) can introduce extra leading dimensions during
349
+ # compilation so we need to reshape back to token-flattened shapes
350
+ q_pe = q_pe.reshape(-1, self.n_head, self.rope_dim)
351
+ k_pe = k_pe.reshape(-1, 1, self.rope_dim)
352
+
353
+ # `rotary_emb` is shape-preserving; `q_pe` is already
354
+ # [num_tokens, n_head, rope_dim].
355
+ q = torch.cat([q_pe, q_nope], dim=-1)
356
+ # `k_pe` is [num_tokens, 1, rope_dim] (MQA).
357
+ k = torch.cat([k_pe.squeeze(-2), k_nope], dim=-1)
358
+ # else: qk_rope_head_dim=0 — no rope component. q is already
359
+ # [num_tokens, n_head, head_dim] and k is [num_tokens, head_dim] (all
360
+ # nope), so skip the rope split / rotary / cat entirely; otherwise the
361
+ # split/reshape would build 0-element tensors (breaks dynamo tracing).
362
+
363
+ # Rotate Q into the cached K basis before computing fp8 MQA logits.
364
+ # Fusing the fp32 FWHT and quantization avoids an intermediate HBM
365
+ # round-trip and bf16 matrix-rounding bias.
366
+ assert self.head_dim == 128 and self.quant_block_size == 128
367
+ assert self.scale_fmt == "ue8m0"
368
+ q = q.view(-1, self.head_dim)
369
+ q_fp8, q_scale = fwht128_quant_fp8(q)
370
+ q_fp8 = q_fp8.view(-1, self.n_head, self.head_dim)
371
+ q_scale = q_scale.view(-1, self.n_head, 1)
372
+
373
+ weights = _fused_indexer_weight_scale(
374
+ weights, q_scale, self.softmax_scale * self.n_head**-0.5
375
+ )
376
+
377
+ # kpool: per-token gate score driving the softmax-weighted pool. Computed
378
+ # from the same hidden_states that produced `k`, so it stays token-aligned.
379
+ # F.linear(x, gate) = x @ gate.T with gate [head_dim, hidden_size].
380
+ gate_score = F.linear(hidden_states, self.index_kpool_compress_gate)
381
+
382
+ # DeepGEMM's MQA-logits kernels (fp8_mqa_logits /
383
+ # fp8_fp4_paged_mqa_logits) require num_heads in {32, 64}; this
384
+ # checkpoint uses index_n_heads=16. Zero-pad q and the per-head
385
+ # weights: logits are a weights-weighted sum over heads, so
386
+ # zero-weight padded heads contribute exactly nothing.
387
+ if self.n_head < 32:
388
+ pad = 32 - self.n_head
389
+ q_fp8 = _pad_indexer_heads(q_fp8, pad)
390
+ weights = _pad_indexer_heads(weights, pad)
391
+
392
+ return self.indexer_op(
393
+ hidden_states,
394
+ q_fp8,
395
+ k,
396
+ weights,
397
+ gate_score=gate_score,
398
+ compress_ape=self.index_kpool_compress_ape,
399
+ index_kpool=self.index_kpool,
400
+ positions=positions,
401
+ )
402
+
403
+
404
+class Glm5NextMLAAttention(nn.Module):
405
+ def __init__(
406
+ self,
407
+ vllm_config: VllmConfig,
408
+ config: Glm5NextConfig,
409
+ hidden_size: int,
410
+ num_heads: int,
411
+ qk_nope_head_dim: int,
412
+ qk_rope_head_dim: int,
413
+ v_head_dim: int,
414
+ q_lora_rank: int | None,
415
+ kv_lora_rank: int,
416
+ max_position_embeddings: int = 8192,
417
+ cache_config: CacheConfig | None = None,
418
+ quant_config: QuantizationConfig | None = None,
419
+ prefix: str = "",
420
+ topk_indices_buffer: torch.Tensor | None = None,
421
+ input_size: int | None = None,
422
+ skip_rope: bool | None = False,
423
+ ) -> None:
424
+ super().__init__()
425
+ self.hidden_size = hidden_size
426
+ self.qk_nope_head_dim = qk_nope_head_dim
427
+ self.qk_rope_head_dim = qk_rope_head_dim
428
+ self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
429
+ self.v_head_dim = v_head_dim
430
+
431
+ self.q_lora_rank = q_lora_rank
432
+ self.kv_lora_rank = kv_lora_rank
433
+
434
+ self.num_heads = num_heads
435
+ tp_size = get_tensor_model_parallel_world_size()
436
+ assert num_heads % tp_size == 0
437
+ self.num_local_heads = num_heads // tp_size
438
+
439
+ self.scaling = self.qk_head_dim**-0.5
440
+ self.max_position_embeddings = max_position_embeddings
441
+
442
+ # Use input_size for projection input dimensions if provided,
443
+ # otherwise default to hidden_size (used in Eagle3 Deepseek with MLA)
444
+ proj_input_size = input_size if input_size is not None else self.hidden_size
445
+
446
+ if self.q_lora_rank is not None:
447
+ self.fused_qkv_a_proj = DeepSeekV2FusedQkvAProjLinear(
448
+ proj_input_size,
449
+ [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
450
+ quant_config=quant_config,
451
+ prefix=f"{prefix}.fused_qkv_a_proj",
452
+ )
453
+ else:
454
+ self.kv_a_proj_with_mqa = ReplicatedLinear(
455
+ proj_input_size,
456
+ self.kv_lora_rank + self.qk_rope_head_dim,
457
+ bias=False,
458
+ quant_config=quant_config,
459
+ prefix=f"{prefix}.kv_a_proj_with_mqa",
460
+ )
461
+
462
+ if self.q_lora_rank is not None:
463
+ self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps)
464
+ self.q_b_proj = ColumnParallelLinear(
465
+ self.q_lora_rank,
466
+ self.num_heads * self.qk_head_dim,
467
+ bias=False,
468
+ quant_config=quant_config,
469
+ prefix=f"{prefix}.q_b_proj",
470
+ )
471
+ else:
472
+ self.q_proj = ColumnParallelLinear(
473
+ proj_input_size,
474
+ self.num_heads * self.qk_head_dim,
475
+ bias=False,
476
+ quant_config=quant_config,
477
+ prefix=f"{prefix}.q_proj",
478
+ )
479
+ self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps)
480
+ self.kv_b_proj = ColumnParallelLinear(
481
+ self.kv_lora_rank,
482
+ self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
483
+ bias=False,
484
+ quant_config=quant_config,
485
+ prefix=f"{prefix}.kv_b_proj",
486
+ )
487
+ self.o_proj = RowParallelLinear(
488
+ self.num_heads * self.v_head_dim,
489
+ self.hidden_size,
490
+ bias=False,
491
+ quant_config=quant_config,
492
+ prefix=f"{prefix}.o_proj",
493
+ )
494
+
495
+ if not skip_rope:
496
+ assert config.rope_parameters is not None
497
+ if config.rope_parameters["rope_type"] != "default":
498
+ config.rope_parameters["rope_type"] = (
499
+ "deepseek_yarn"
500
+ if config.rope_parameters.get("apply_yarn_scaling", True)
501
+ else "deepseek_llama_scaling"
502
+ )
503
+
504
+ self.rotary_emb: RotaryEmbedding | None = get_rope(
505
+ qk_rope_head_dim,
506
+ max_position=max_position_embeddings,
507
+ rope_parameters=config.rope_parameters,
508
+ is_neox_style=False,
509
+ )
510
+
511
+ if (
512
+ config.rope_parameters["rope_type"] != "default"
513
+ and config.rope_parameters["rope_type"] == "deepseek_yarn"
514
+ ):
515
+ mscale_all_dim = config.rope_parameters.get("mscale_all_dim", False)
516
+ scaling_factor = config.rope_parameters["factor"]
517
+ mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim))
518
+ self.scaling = self.scaling * mscale * mscale
519
+ else:
520
+ self.rotary_emb = None
521
+
522
+ self.is_v32 = config.index_topk is not None
523
+
524
+ if self.is_v32:
525
+ self.indexer_rope_emb: RotaryEmbedding | None = get_rope(
526
+ qk_rope_head_dim,
527
+ max_position=max_position_embeddings,
528
+ rope_parameters=config.rope_parameters,
529
+ is_neox_style=not config.indexer_rope_interleave,
530
+ )
531
+ # The sparse indexer projects from the MLA q-lora rank, which is
532
+ # always set for v32 MLA configs; narrow away the `int | None`.
533
+ assert q_lora_rank is not None
534
+ self.indexer: Indexer | None = Indexer(
535
+ vllm_config,
536
+ config,
537
+ hidden_size,
538
+ q_lora_rank,
539
+ quant_config,
540
+ cache_config,
541
+ topk_indices_buffer,
542
+ f"{prefix}.indexer",
543
+ )
544
+
545
+ else:
546
+ self.indexer_rope_emb = None
547
+ self.indexer = None
548
+
549
+ mla_modules = MLAModules(
550
+ kv_a_layernorm=self.kv_a_layernorm,
551
+ kv_b_proj=self.kv_b_proj,
552
+ rotary_emb=self.rotary_emb,
553
+ o_proj=self.o_proj,
554
+ fused_qkv_a_proj=self.fused_qkv_a_proj
555
+ if self.q_lora_rank is not None
556
+ else None,
557
+ kv_a_proj_with_mqa=self.kv_a_proj_with_mqa
558
+ if self.q_lora_rank is None
559
+ else None,
560
+ q_a_layernorm=self.q_a_layernorm if self.q_lora_rank is not None else None,
561
+ q_b_proj=self.q_b_proj if self.q_lora_rank is not None else None,
562
+ q_proj=self.q_proj if self.q_lora_rank is None else None,
563
+ indexer=self.indexer,
564
+ indexer_rotary_emb=self.indexer_rope_emb,
565
+ is_sparse=self.is_v32,
566
+ topk_indices_buffer=topk_indices_buffer,
567
+ )
568
+
569
+ self.mla_attn = MultiHeadLatentAttentionWrapper(
570
+ self.hidden_size,
571
+ self.num_local_heads,
572
+ self.scaling,
573
+ self.qk_nope_head_dim,
574
+ self.qk_rope_head_dim,
575
+ self.v_head_dim,
576
+ self.q_lora_rank,
577
+ self.kv_lora_rank,
578
+ mla_modules,
579
+ cache_config,
580
+ quant_config,
581
+ prefix,
582
+ skip_topk=False,
583
+ fuse_qkv_rmsnorm=True,
584
+ )
585
+
586
+ def forward(
587
+ self, hidden_states: torch.Tensor, positions: torch.Tensor
588
+ ) -> torch.Tensor:
589
+ # The wrapper also runs the sparse indexer before MLA attention.
590
+ return self.mla_attn(positions, hidden_states)
vllm/models/glm5next/nvidia/kda.py
new
+611
@@ -0,0 +1,611 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""GLM-5.3-Flash KDA layer with separate convolutions and a bounded safe gate."""
4
+
5
+import torch
6
+from torch import nn
7
+
8
+from vllm.compilation.breakable_cudagraph import eager_break_during_capture
9
+from vllm.config import VllmConfig, get_current_vllm_config
10
+from vllm.distributed import divide
11
+from vllm.forward_context import get_forward_context
12
+from vllm.model_executor.layers.linear import (
13
+ ColumnParallelLinear,
14
+ MergedColumnParallelLinear,
15
+ RowParallelLinear,
16
+)
17
+from vllm.model_executor.layers.mamba.gdn.base import GatedDeltaNetAttention
18
+from vllm.model_executor.layers.mamba.mamba_utils import (
19
+ MambaStateDtypeCalculator,
20
+ MambaStateShapeCalculator,
21
+ is_conv_state_dim_first,
22
+)
23
+from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
24
+ causal_conv1d_fn,
25
+ causal_conv1d_update,
26
+)
27
+from vllm.model_executor.layers.mamba.ops.gather_initial_states import (
28
+ gather_initial_states,
29
+)
30
+from vllm.model_executor.layers.mamba.ops.scatter_states import scatter_states
31
+from vllm.model_executor.model_loader.weight_utils import sharded_weight_loader
32
+from vllm.model_executor.utils import (
33
+ maybe_disable_graph_partition,
34
+ set_weight_attrs,
35
+)
36
+from vllm.platforms import current_platform
37
+from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated
38
+from vllm.transformers_utils.configs.glm5_next import Glm5NextConfig
39
+from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
40
+
41
+if current_platform.is_rocm():
42
+ from vllm.models.glm5next.amd.ops.third_party.kda import (
43
+ chunk_kda_with_fused_gate,
44
+ fused_recurrent_kda,
45
+ )
46
+else:
47
+ from vllm.models.glm5next.nvidia.ops.third_party.kda import (
48
+ chunk_kda_with_fused_gate,
49
+ fused_recurrent_kda,
50
+ )
51
+
52
+
53
+class _Glm5NextMergedColumnParallelLinear(MergedColumnParallelLinear):
54
+ """Merged projection with multiple replicated output shards.
55
+
56
+ Extends K3's ``_KimiGDNMergedColumnParallelLinear`` to support two
57
+ replicated shards (f_a, g_a) instead of one. Pre-multiplies each
58
+ replicated entry's output_size by tp_size so the per-rank shard
59
+ divides back to the full size, and forces tp_rank=0 during weight
60
+ loading for replicated shards.
61
+ """
62
+
63
+ def __init__(
64
+ self,
65
+ input_size: int,
66
+ output_sizes: list[int],
67
+ replicated_shard_ids: tuple[int, ...],
68
+ tp_size: int,
69
+ **kwargs,
70
+ ) -> None:
71
+ self.replicated_shard_ids = set(replicated_shard_ids)
72
+ output_sizes = output_sizes.copy()
73
+ for sid in self.replicated_shard_ids:
74
+ output_sizes[sid] *= tp_size
75
+ super().__init__(input_size, output_sizes, **kwargs)
76
+
77
+ def weight_loader(
78
+ self,
79
+ param: nn.Parameter,
80
+ loaded_weight: torch.Tensor,
81
+ loaded_shard_id: tuple[int, ...] | int | None = None,
82
+ ) -> None:
83
+ tp_rank = self.tp_rank
84
+ param_tp_rank = getattr(param, "tp_rank", None)
85
+ if loaded_shard_id in self.replicated_shard_ids:
86
+ self.tp_rank = 0
87
+ if param_tp_rank is not None:
88
+ param.tp_rank = 0
89
+ try:
90
+ super().weight_loader(param, loaded_weight, loaded_shard_id)
91
+ finally:
92
+ self.tp_rank = tp_rank
93
+ if param_tp_rank is not None:
94
+ param.tp_rank = param_tp_rank
95
+
96
+ def weight_loader_v2(
97
+ self,
98
+ param: nn.Parameter,
99
+ loaded_weight: torch.Tensor,
100
+ loaded_shard_id: tuple[int, ...] | int | None = None,
101
+ ) -> None:
102
+ tp_rank = self.tp_rank
103
+ param_tp_rank = getattr(param, "tp_rank", None)
104
+ if loaded_shard_id in self.replicated_shard_ids:
105
+ self.tp_rank = 0
106
+ if param_tp_rank is not None:
107
+ param.tp_rank = 0
108
+ try:
109
+ super().weight_loader_v2(param, loaded_weight, loaded_shard_id)
110
+ finally:
111
+ self.tp_rank = tp_rank
112
+ if param_tp_rank is not None:
113
+ param.tp_rank = param_tp_rank
114
+
115
+
116
+@torch.compile(
117
+ dynamic=True,
118
+ backend=current_platform.simple_compile_backend,
119
+ options=maybe_disable_graph_partition(current_platform.simple_compile_backend),
120
+)
121
+def _cast_sigmoid(x: torch.Tensor) -> torch.Tensor:
122
+ """Fuse the fp32 cast + sigmoid into one Inductor kernel."""
123
+ return x.float().sigmoid()
124
+
125
+
126
+class Glm5NextLinearAttention(GatedDeltaNetAttention):
127
+ head_dim: int
128
+ num_heads: int
129
+ conv_size: int
130
+
131
+ def get_state_dtype(
132
+ self,
133
+ ) -> tuple[torch.dtype, torch.dtype]:
134
+ if self.model_config is None or self.cache_config is None:
135
+ raise ValueError("model_config and cache_config must be set")
136
+ return MambaStateDtypeCalculator.kda_state_dtype(
137
+ self.model_config.dtype, self.cache_config.mamba_cache_dtype
138
+ )
139
+
140
+ def get_state_shape(
141
+ self,
142
+ ) -> tuple[tuple[int, ...], tuple[int, ...]]:
143
+ # conv_state width must include num_spec so the spec-decode conv update
144
+ # (causal_conv1d_update with num_accepted_tokens + max_query_len) can
145
+ # slide the window across the draft-verify tokens without reading past
146
+ # the allocated width. Matches qwen_gdn_linear_attn.get_state_shape.
147
+ return MambaStateShapeCalculator.kda_state_shape(
148
+ self.tp_size,
149
+ self.num_heads,
150
+ self.head_dim,
151
+ conv_kernel_size=self.conv_size,
152
+ num_spec=self.num_spec,
153
+ )
154
+
155
+ def __init__(
156
+ self,
157
+ config: Glm5NextConfig,
158
+ vllm_config: VllmConfig,
159
+ prefix: str = "",
160
+ ) -> None:
161
+ # KDA projections remain BF16 because fp8 checkpoints omit their scales.
162
+ saved_quant_config = vllm_config.quant_config
163
+ try:
164
+ vllm_config.quant_config = None
165
+ super().__init__(config, vllm_config, prefix)
166
+ finally:
167
+ vllm_config.quant_config = saved_quant_config
168
+
169
+ self.head_dim = config.linear_head_dim
170
+ self.num_heads = config.linear_num_heads
171
+ self.conv_size = config.linear_conv_kernel_dim
172
+ assert self.num_heads % self.tp_size == 0
173
+ self.local_num_heads = divide(self.num_heads, self.tp_size)
174
+
175
+ projection_size = self.head_dim * self.num_heads
176
+ self.local_projection_size = divide(projection_size, self.tp_size)
177
+
178
+ # Merge q, k, v, b, f_a, g_a projections into one GEMM (6→1 launches).
179
+ # Order matches checkpoint's fused_qkvbfg_a_proj convention.
180
+ # Shards 4 (f_a) and 5 (g_a) are replicated across TP ranks.
181
+ self.in_proj_qkvbfg_a = _Glm5NextMergedColumnParallelLinear(
182
+ self.hidden_size,
183
+ [
184
+ projection_size, # q (shard 0)
185
+ projection_size, # k (shard 1)
186
+ projection_size, # v (shard 2)
187
+ self.num_heads, # b (shard 3)
188
+ self.head_dim, # f_a (shard 4, replicated)
189
+ self.head_dim, # g_a (shard 5, replicated)
190
+ ],
191
+ replicated_shard_ids=(4, 5),
192
+ tp_size=self.tp_size,
193
+ bias=False,
194
+ quant_config=self.quant_config,
195
+ prefix=f"{prefix}.in_proj_qkvbfg_a",
196
+ )
197
+
198
+ self.f_b_proj = ColumnParallelLinear(
199
+ self.head_dim,
200
+ projection_size,
201
+ bias=False,
202
+ quant_config=self.quant_config,
203
+ prefix=f"{prefix}.f_b_proj",
204
+ )
205
+ self.dt_bias = nn.Parameter(
206
+ torch.empty(divide(projection_size, self.tp_size), dtype=torch.float32)
207
+ )
208
+
209
+ set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)})
210
+
211
+ self.q_conv1d = ColumnParallelLinear(
212
+ input_size=self.conv_size,
213
+ output_size=projection_size,
214
+ bias=False,
215
+ params_dtype=torch.float32,
216
+ prefix=f"{prefix}.q_conv1d",
217
+ )
218
+ self.k_conv1d = ColumnParallelLinear(
219
+ input_size=self.conv_size,
220
+ output_size=projection_size,
221
+ bias=False,
222
+ params_dtype=torch.float32,
223
+ prefix=f"{prefix}.k_conv1d",
224
+ )
225
+ self.v_conv1d = ColumnParallelLinear(
226
+ input_size=self.conv_size,
227
+ output_size=projection_size,
228
+ bias=False,
229
+ params_dtype=torch.float32,
230
+ prefix=f"{prefix}.v_conv1d",
231
+ )
232
+ # unsqueeze to fit conv1d weights shape into the linear weights shape.
233
+ # Can't do this in `weight_loader` since it already exists in
234
+ # `ColumnParallelLinear` and `set_weight_attrs`
235
+ # doesn't allow to override it
236
+ self.q_conv1d.weight.data = self.q_conv1d.weight.data.unsqueeze(1)
237
+ self.k_conv1d.weight.data = self.k_conv1d.weight.data.unsqueeze(1)
238
+ self.v_conv1d.weight.data = self.v_conv1d.weight.data.unsqueeze(1)
239
+ # Lazily-built merged q|k|v conv weight (built on first forward, after
240
+ # weights are loaded). See _forward.
241
+ self._merged_conv_weight: torch.Tensor | None = None
242
+
243
+ self.A_log = nn.Parameter(
244
+ torch.empty(1, 1, self.local_num_heads, 1, dtype=torch.float32)
245
+ )
246
+ set_weight_attrs(self.A_log, {"weight_loader": sharded_weight_loader(2)})
247
+
248
+ self.g_b_proj = ColumnParallelLinear(
249
+ self.head_dim,
250
+ projection_size,
251
+ bias=False,
252
+ quant_config=self.quant_config,
253
+ prefix=f"{prefix}.g_b_proj",
254
+ )
255
+ self.o_norm = FusedRMSNormGated(self.head_dim, activation="sigmoid")
256
+ self.o_proj = RowParallelLinear(
257
+ projection_size,
258
+ self.hidden_size,
259
+ bias=False,
260
+ quant_config=self.quant_config,
261
+ prefix=f"{prefix}.o_proj",
262
+ )
263
+
264
+ compilation_config = get_current_vllm_config().compilation_config
265
+ if prefix in compilation_config.static_forward_context:
266
+ raise ValueError(f"Duplicate layer name: {prefix}")
267
+ compilation_config.static_forward_context[prefix] = self
268
+
269
+ # Checkpoints store A_log as 1-D; the model parameter is 4-D.
270
+ def _a_log_weight_loader(param, loaded_weight):
271
+ if loaded_weight.dim() == 1:
272
+ loaded_weight = loaded_weight.view([1, 1, -1, 1])
273
+ return sharded_weight_loader(2)(param, loaded_weight)
274
+
275
+ self.A_log.weight_loader = _a_log_weight_loader
276
+
277
+ # GLM-5.3-Flash uses a bounded sigmoid gate instead of the default
278
+ # unbounded softplus gate.
279
+ self.kda_safe_gate = True
280
+ self.kda_lower_bound = config.linear_lower_bound
281
+ # Process-global conv-state layout, resolved once here instead of on
282
+ # every _forward call (it reads an env-derived flag each time).
283
+ self._conv_state_dim_first = is_conv_state_dim_first()
284
+
285
+ def forward(
286
+ self,
287
+ hidden_states: torch.Tensor,
288
+ positions: torch.Tensor,
289
+ ) -> torch.Tensor:
290
+ num_tokens = hidden_states.size(0)
291
+ # One merged GEMM for q, k, v, b, f_a, g_a (replaces 6 separate GEMMs).
292
+ projected = self.in_proj_qkvbfg_a(hidden_states)[0]
293
+ qkv, beta_raw, f_a, g_a = projected.split(
294
+ [
295
+ 3 * self.local_projection_size,
296
+ self.local_num_heads,
297
+ self.head_dim,
298
+ self.head_dim,
299
+ ],
300
+ dim=-1,
301
+ )
302
+
303
+ # Beta stays raw (bf16) here: the recurrent kernel sigmoids it in fp32
304
+ # at load (SIGMOID_BETA), and only the chunked prefill path needs the
305
+ # pre-computed fp32 sigmoid — computed lazily in _forward. Pure decode
306
+ # / spec-verify steps then skip the _cast_sigmoid kernel and its fp32
307
+ # intermediate entirely.
308
+ beta = beta_raw.unsqueeze(0)
309
+ g1 = self.f_b_proj(f_a)[0]
310
+ g1 = g1.reshape(1, -1, self.local_num_heads, self.head_dim)
311
+
312
+ g_proj_states = self.g_b_proj(g_a)[0]
313
+ # Must stay 3D: rms_norm_gated reads H from g.shape[-2].
314
+ g2 = g_proj_states.reshape(-1, self.local_num_heads, self.head_dim)
315
+
316
+ core_attn_out = torch.empty(
317
+ (1, num_tokens, self.local_num_heads, self.head_dim),
318
+ dtype=hidden_states.dtype,
319
+ device=hidden_states.device,
320
+ )
321
+ # Call the decorated eager break directly so host-side prefill branches
322
+ # are not captured by PIECEWISE CUDA graphs.
323
+ self._forward(
324
+ qkv_proj_states=qkv,
325
+ g1=g1,
326
+ beta=beta,
327
+ core_attn_out=core_attn_out,
328
+ )
329
+ core_attn_out = self.o_norm(core_attn_out, g2)
330
+ core_attn_out = core_attn_out.reshape(core_attn_out.size(1), -1)
331
+ return self.o_proj(core_attn_out)[0]
332
+
333
+ @eager_break_during_capture
334
+ def _forward(
335
+ self,
336
+ qkv_proj_states: torch.Tensor,
337
+ g1: torch.Tensor,
338
+ beta: torch.Tensor,
339
+ core_attn_out: torch.Tensor,
340
+ ) -> None:
341
+ forward_context = get_forward_context()
342
+ attn_metadata_raw = forward_context.attn_metadata
343
+
344
+ if attn_metadata_raw is None:
345
+ return
346
+
347
+ assert isinstance(attn_metadata_raw, dict)
348
+ attn_metadata_narrowed = attn_metadata_raw.get(self.prefix)
349
+ if attn_metadata_narrowed is None:
350
+ # Profile/warmup dummy runs may omit mamba-family metadata.
351
+ return
352
+ assert isinstance(attn_metadata_narrowed, GDNAttentionMetadata)
353
+ has_initial_state = attn_metadata_narrowed.has_initial_state
354
+ non_spec_query_start_loc = attn_metadata_narrowed.non_spec_query_start_loc
355
+ non_spec_state_indices_tensor = (
356
+ attn_metadata_narrowed.non_spec_state_indices_tensor
357
+ ) # noqa: E501
358
+ num_actual_tokens = attn_metadata_narrowed.num_actual_tokens
359
+ # Spec-decode metadata (all None when speculative decoding is disabled).
360
+ spec_sequence_masks = attn_metadata_narrowed.spec_sequence_masks
361
+ spec_query_start_loc = attn_metadata_narrowed.spec_query_start_loc
362
+ spec_state_indices_tensor = attn_metadata_narrowed.spec_state_indices_tensor
363
+ spec_token_indx = attn_metadata_narrowed.spec_token_indx
364
+ non_spec_token_indx = attn_metadata_narrowed.non_spec_token_indx
365
+ num_accepted_tokens = attn_metadata_narrowed.num_accepted_tokens
366
+ num_spec_decodes = attn_metadata_narrowed.num_spec_decodes
367
+ use_spec = spec_sequence_masks is not None and num_spec_decodes > 0
368
+ # Safe-gate checkpoints use the bounded sigmoid variant.
369
+ safe_gate = self.kda_safe_gate
370
+ lower_bound = self.kda_lower_bound
371
+ constant_caches = self.kv_cache
372
+
373
+ qkv_proj_states = qkv_proj_states[:num_actual_tokens]
374
+ g1 = g1[:, :num_actual_tokens]
375
+ beta = beta[:, :num_actual_tokens]
376
+
377
+ (conv_state, recurrent_state) = constant_caches
378
+ # conv_state must be (..., dim, width-1) for the conv kernels.
379
+ # DS layout stores it that way directly; SD layout needs a transpose.
380
+ # Layout is process-global and resolved once at init (see __init__).
381
+ if not self._conv_state_dim_first:
382
+ conv_state = conv_state.transpose(-1, -2)
383
+
384
+ # One merged short-conv over q|k|v instead of three separate calls. The
385
+ # 1D conv is independent per channel, so concatenating q/k/v along the
386
+ # channel dim and running a single causal_conv1d is bit-identical to
387
+ # three calls. The merged weight is q|k|v conv weights concatenated;
388
+ # built once and cached (params are fixed after load). conv_state is
389
+ # already stored as the merged q|k|v state, so it is used directly.
390
+ if self._merged_conv_weight is None:
391
+
392
+ def _w(m):
393
+ return m.weight.view(m.weight.size(0), m.weight.size(2))
394
+
395
+ self._merged_conv_weight = torch.cat(
396
+ [_w(self.q_conv1d), _w(self.k_conv1d), _w(self.v_conv1d)],
397
+ dim=0,
398
+ ).contiguous()
399
+ conv_weights = self._merged_conv_weight
400
+ conv_bias = self.q_conv1d.bias
401
+
402
+ # Split projections / gating into spec (draft-verify) and non-spec token
403
+ # groups when speculative decoding is active. Spec tokens carry
404
+ # num_spec+1 recurrent-state columns each and are advanced with
405
+ # num_accepted_tokens for rejection-sampling rollback; non-spec tokens
406
+ # are one-per-request. Mirrors olmo_gdn_linear_attn.py. Projections are
407
+ # [n, *] (token dim 0); g1/beta are [1, n, h, d] (token dim 1).
408
+ if use_spec:
409
+ # In a pure spec-verify step (no non-spec tokens) the metadata
410
+ # builder sets spec_token_indx = arange(num_actual_tokens), making
411
+ # the index_select calls below identity copies. Skip them on this
412
+ # steady-state decode hot path. The outputs alias the inputs here;
413
+ # the downstream conv/recurrent kernels read them without mutating
414
+ # in place, so the aliasing is safe.
415
+ if non_spec_token_indx is None or non_spec_token_indx.numel() == 0:
416
+ qkv_spec = qkv_proj_states
417
+ g1_spec = g1
418
+ beta_spec = beta
419
+ else:
420
+ qkv_spec = qkv_proj_states.index_select(0, spec_token_indx)
421
+ g1_spec = g1.index_select(1, spec_token_indx)
422
+ beta_spec = beta.index_select(1, spec_token_indx)
423
+ if non_spec_token_indx is not None and non_spec_token_indx.numel() > 0:
424
+ qkv_ns = qkv_proj_states.index_select(0, non_spec_token_indx)
425
+ g1_ns = g1.index_select(1, non_spec_token_indx)
426
+ beta_ns = beta.index_select(1, non_spec_token_indx)
427
+ else:
428
+ qkv_ns = g1_ns = beta_ns = None
429
+ else:
430
+ qkv_spec = g1_spec = beta_spec = None
431
+ qkv_ns, g1_ns, beta_ns = qkv_proj_states, g1, beta
432
+
433
+ # --- causal conv1d: spec (draft-verify) path ---
434
+ if use_spec:
435
+ assert spec_state_indices_tensor is not None
436
+ assert num_accepted_tokens is not None
437
+ conv_idx = spec_state_indices_tensor[:, 0][:num_spec_decodes]
438
+ conv_mql = spec_state_indices_tensor.size(-1)
439
+ qkv_spec = causal_conv1d_update(
440
+ qkv_spec,
441
+ conv_state,
442
+ conv_weights,
443
+ conv_bias,
444
+ activation="silu",
445
+ conv_state_indices=conv_idx,
446
+ num_accepted_tokens=num_accepted_tokens,
447
+ query_start_loc=spec_query_start_loc,
448
+ max_query_len=conv_mql,
449
+ )
450
+ q_spec, k_spec, v_spec = qkv_spec.split(self.local_projection_size, dim=-1)
451
+
452
+ # --- causal conv1d: non-spec path (prefill or plain decode) ---
453
+ q_ns = k_ns = v_ns = None
454
+ if attn_metadata_narrowed.num_prefills > 0:
455
+ assert qkv_ns is not None
456
+ qkv_ns = causal_conv1d_fn(
457
+ qkv_ns.transpose(0, 1),
458
+ conv_weights,
459
+ conv_bias,
460
+ activation="silu",
461
+ conv_states=conv_state,
462
+ has_initial_state=has_initial_state,
463
+ cache_indices=non_spec_state_indices_tensor,
464
+ query_start_loc=non_spec_query_start_loc,
465
+ metadata=attn_metadata_narrowed,
466
+ ).transpose(0, 1)
467
+ q_ns, k_ns, v_ns = qkv_ns.split(self.local_projection_size, dim=-1)
468
+ elif attn_metadata_narrowed.num_decodes > 0:
469
+ assert non_spec_state_indices_tensor is not None
470
+ decode_conv_indices = non_spec_state_indices_tensor[
471
+ : attn_metadata_narrowed.num_decodes
472
+ ]
473
+ qkv_ns = causal_conv1d_update(
474
+ qkv_ns,
475
+ conv_state,
476
+ conv_weights,
477
+ conv_bias,
478
+ activation="silu",
479
+ conv_state_indices=decode_conv_indices,
480
+ )
481
+ q_ns, k_ns, v_ns = qkv_ns.split(self.local_projection_size, dim=-1)
482
+
483
+ def _rearr(x):
484
+ return x.reshape(1, -1, self.local_num_heads, self.head_dim)
485
+
486
+ # --- core attention: spec (draft-verify) path ---
487
+ core_attn_out_spec = None
488
+ # In a pure spec-verify step (no non-spec tokens) the recurrent kernel
489
+ # can write straight into the layer output buffer, skipping the
490
+ # fresh allocation + copy below. Mixed steps must scatter via
491
+ # spec_token_indx, so they keep the kernel-managed output.
492
+ spec_out = (
493
+ core_attn_out[0, :num_actual_tokens].unsqueeze(0)
494
+ if non_spec_token_indx is None or non_spec_token_indx.numel() == 0
495
+ else None
496
+ )
497
+ if use_spec:
498
+ assert spec_state_indices_tensor is not None
499
+ assert num_accepted_tokens is not None
500
+ assert spec_query_start_loc is not None
501
+ # Gate computed inside the recurrent kernel (COMPUTE_GATE) from
502
+ # raw g1 — replicates fused_kda_gate's arithmetic bit-for-bit and
503
+ # skips its launch + fp32 [n, H, D] intermediate per layer.
504
+ core_attn_out_spec, _ = fused_recurrent_kda(
505
+ q=_rearr(q_spec),
506
+ k=_rearr(k_spec),
507
+ v=_rearr(v_spec),
508
+ g=g1_spec,
509
+ beta=beta_spec,
510
+ initial_state=recurrent_state,
511
+ use_qk_l2norm_in_kernel=True,
512
+ cu_seqlens=spec_query_start_loc[: num_spec_decodes + 1],
513
+ ssm_state_indices=spec_state_indices_tensor,
514
+ num_accepted_tokens=num_accepted_tokens,
515
+ out=spec_out,
516
+ sigmoid_beta=True,
517
+ a_log=self.A_log,
518
+ g_bias=self.dt_bias,
519
+ compute_gate=True,
520
+ lower_bound=lower_bound,
521
+ )
522
+
523
+ # --- core attention: non-spec path (prefill or plain decode) ---
524
+ core_attn_out_non_spec = None
525
+ # Only the plain-decode recurrent kernel can write straight into the
526
+ # layer output buffer; the chunked prefill kernel cannot, so this
527
+ # stays None there and the merge copy below runs as before.
528
+ ns_out = None
529
+ if attn_metadata_narrowed.num_prefills > 0:
530
+ assert q_ns is not None
531
+ assert non_spec_state_indices_tensor is not None
532
+ assert has_initial_state is not None
533
+ initial_state = gather_initial_states(
534
+ recurrent_state, non_spec_state_indices_tensor, has_initial_state
535
+ )
536
+ (
537
+ core_attn_out_non_spec,
538
+ last_recurrent_state,
539
+ ) = chunk_kda_with_fused_gate(
540
+ q=_rearr(q_ns),
541
+ k=_rearr(k_ns),
542
+ v=_rearr(v_ns),
543
+ raw_g=g1_ns,
544
+ # Chunk path wants the pre-sigmoided fp32 beta (its kernels
545
+ # don't sigmoid); beta_ns is raw bf16 from forward.
546
+ beta=_cast_sigmoid(beta_ns.squeeze(0)).unsqueeze(0),
547
+ A_log=self.A_log,
548
+ g_bias=self.dt_bias,
549
+ initial_state=initial_state,
550
+ output_final_state=True,
551
+ use_qk_l2norm_in_kernel=True,
552
+ cu_seqlens=non_spec_query_start_loc,
553
+ safe_gate=safe_gate,
554
+ lower_bound=lower_bound,
555
+ )
556
+ # Init cache
557
+ scatter_states(
558
+ recurrent_state,
559
+ last_recurrent_state,
560
+ non_spec_state_indices_tensor,
561
+ )
562
+ elif attn_metadata_narrowed.num_decodes > 0:
563
+ assert non_spec_query_start_loc is not None
564
+ assert non_spec_state_indices_tensor is not None
565
+ # Plain decode step (no spec tokens): token order is dense, so the
566
+ # kernel can write straight into the layer output buffer. A mixed
567
+ # step scatters non-spec output via non_spec_token_indx instead.
568
+ # Gate computed in-kernel (COMPUTE_GATE), beta sigmoided in-kernel.
569
+ if not use_spec:
570
+ ns_out = spec_out
571
+ core_attn_out_non_spec, _ = fused_recurrent_kda(
572
+ q=_rearr(q_ns),
573
+ k=_rearr(k_ns),
574
+ v=_rearr(v_ns),
575
+ g=g1_ns,
576
+ beta=beta_ns,
577
+ initial_state=recurrent_state,
578
+ use_qk_l2norm_in_kernel=True,
579
+ cu_seqlens=non_spec_query_start_loc[
580
+ : attn_metadata_narrowed.num_decodes + 1
581
+ ],
582
+ ssm_state_indices=non_spec_state_indices_tensor,
583
+ out=ns_out,
584
+ sigmoid_beta=True,
585
+ a_log=self.A_log,
586
+ g_bias=self.dt_bias,
587
+ compute_gate=True,
588
+ lower_bound=lower_bound,
589
+ )
590
+
591
+ # --- merge spec / non-spec outputs back into token order ---
592
+ if use_spec and core_attn_out_non_spec is not None:
593
+ assert core_attn_out_spec is not None
594
+ merged = torch.empty(
595
+ (1, num_actual_tokens, *core_attn_out_spec.shape[2:]),
596
+ dtype=core_attn_out_non_spec.dtype,
597
+ device=core_attn_out_non_spec.device,
598
+ )
599
+ merged.index_copy_(1, spec_token_indx, core_attn_out_spec)
600
+ merged.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec)
601
+ core_attn_out[0, :num_actual_tokens] = merged.squeeze(0)
602
+ elif use_spec:
603
+ assert core_attn_out_spec is not None
604
+ if spec_out is None:
605
+ core_attn_out[0, :num_actual_tokens] = core_attn_out_spec.squeeze(0)
606
+ else:
607
+ assert core_attn_out_non_spec is not None
608
+ if ns_out is None:
609
+ core_attn_out[0, :num_actual_tokens] = core_attn_out_non_spec[
610
+ 0, :num_actual_tokens
611
+ ]
vllm/models/glm5next/nvidia/model.py
new
+1196
@@ -0,0 +1,1196 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from collections.abc import Iterable
5
+from typing import ClassVar, Literal
6
+
7
+import torch
8
+from torch import nn
9
+
10
+from vllm.config import ParallelConfig, VllmConfig
11
+from vllm.distributed import (
12
+ get_ep_group,
13
+ get_pp_group,
14
+ get_tensor_model_parallel_rank,
15
+ get_tensor_model_parallel_world_size,
16
+ tensor_model_parallel_all_gather,
17
+)
18
+from vllm.logger import init_logger
19
+from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp
20
+from vllm.model_executor.layers.fused_moe import (
21
+ FusedMoEFactory,
22
+ GateLinear,
23
+ fused_moe_make_expert_params_mapping,
24
+)
25
+from vllm.model_executor.layers.layernorm import RMSNorm
26
+from vllm.model_executor.layers.linear import (
27
+ MergedColumnParallelLinear,
28
+ RowParallelLinear,
29
+)
30
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
31
+from vllm.model_executor.layers.mamba.mamba_utils import (
32
+ MambaStateCopyFunc,
33
+ MambaStateCopyFuncCalculator,
34
+ MambaStateDtypeCalculator,
35
+ MambaStateShapeCalculator,
36
+)
37
+from vllm.model_executor.layers.mhc import (
38
+ MHCFusedPostPreOp,
39
+ MHCPostOp,
40
+ MHCPreOp,
41
+ hc_contract,
42
+ hc_expand,
43
+)
44
+from vllm.model_executor.layers.quantization import QuantizationConfig
45
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
46
+ GroupShape,
47
+ scaled_dequantize,
48
+)
49
+from vllm.model_executor.layers.vocab_parallel_embedding import (
50
+ ParallelLMHead,
51
+ VocabParallelEmbedding,
52
+)
53
+from vllm.model_executor.model_loader.weight_utils import (
54
+ default_weight_loader,
55
+ maybe_remap_kv_scale_name,
56
+)
57
+from vllm.model_executor.models.deepseek_v2 import _get_moe_router_dtype
58
+from vllm.model_executor.models.glm4_1v import (
59
+ Glm4vDummyInputsBuilder,
60
+ Glm4vForConditionalGeneration,
61
+)
62
+from vllm.model_executor.models.interfaces import (
63
+ HasInnerState,
64
+ IsHybrid,
65
+ MixtureOfExperts,
66
+ SupportsPP,
67
+)
68
+from vllm.model_executor.models.utils import (
69
+ AutoWeightsLoader,
70
+ PPMissingLayer,
71
+ init_vllm_registered_model,
72
+ is_pp_missing_parameter,
73
+ make_layers,
74
+ maybe_prefix,
75
+ sequence_parallel_chunk,
76
+)
77
+from vllm.models.common.ops.sequence_parallel import (
78
+ sp_all_gather,
79
+ sp_reduce_scatter,
80
+ sp_shard,
81
+)
82
+from vllm.multimodal import MULTIMODAL_REGISTRY
83
+from vllm.platforms import current_platform
84
+from vllm.sequence import IntermediateTensors
85
+from vllm.transformers_utils.configs.glm5_next import Glm5NextConfig
86
+
87
+from .attention import Glm5NextMLAAttention
88
+from .kda import Glm5NextLinearAttention
89
+from .multimodal import (
90
+ Glm5NextMultiModalProcessor,
91
+ Glm5NextProcessingInfo,
92
+ Glm5NextVisionTransformer,
93
+)
94
+
95
+logger = init_logger(__name__)
96
+
97
+
98
+class Glm5NextMLP(nn.Module):
99
+ def __init__(
100
+ self,
101
+ hidden_size: int,
102
+ intermediate_size: int,
103
+ hidden_act: str,
104
+ quant_config: QuantizationConfig | None = None,
105
+ reduce_results: bool = True,
106
+ is_sequence_parallel=False,
107
+ prefix: str = "",
108
+ swiglu_limit: float | None = None,
109
+ ) -> None:
110
+ super().__init__()
111
+
112
+ # If is_sequence_parallel, the input and output tensors are sharded
113
+ # across the ranks within the tp_group. In this case the weights are
114
+ # replicated and no collective ops are needed.
115
+ # Otherwise we use standard TP with an allreduce at the end.
116
+ self.gate_up_proj = MergedColumnParallelLinear(
117
+ hidden_size,
118
+ [intermediate_size] * 2,
119
+ bias=False,
120
+ quant_config=quant_config,
121
+ disable_tp=is_sequence_parallel,
122
+ prefix=f"{prefix}.gate_up_proj",
123
+ )
124
+ self.down_proj = RowParallelLinear(
125
+ intermediate_size,
126
+ hidden_size,
127
+ bias=False,
128
+ quant_config=quant_config,
129
+ reduce_results=reduce_results,
130
+ disable_tp=is_sequence_parallel,
131
+ prefix=f"{prefix}.down_proj",
132
+ )
133
+ if hidden_act != "silu":
134
+ raise ValueError(
135
+ f"Unsupported activation: {hidden_act}. Only silu is supported for now."
136
+ )
137
+
138
+ self.swiglu_limit = swiglu_limit
139
+ if self.swiglu_limit is not None:
140
+ self.act_fn = SiluAndMulWithClamp(swiglu_limit=self.swiglu_limit)
141
+ else:
142
+ self.act_fn = SiluAndMul()
143
+
144
+ def forward(self, x):
145
+ gate_up, _ = self.gate_up_proj(x)
146
+ x = self.act_fn(gate_up)
147
+ x, _ = self.down_proj(x)
148
+ return x
149
+
150
+
151
+class Glm5NextMoE(nn.Module):
152
+ def __init__(
153
+ self,
154
+ config: Glm5NextConfig,
155
+ parallel_config: ParallelConfig,
156
+ quant_config: QuantizationConfig | None = None,
157
+ prefix: str = "",
158
+ apply_routed_scale_to_output: bool = False,
159
+ ):
160
+ super().__init__()
161
+ self.tp_size = get_tensor_model_parallel_world_size()
162
+ self.tp_rank = get_tensor_model_parallel_rank()
163
+
164
+ self.routed_scaling_factor = config.routed_scaling_factor
165
+
166
+ self.ep_group = get_ep_group().device_group
167
+ self.ep_rank = get_ep_group().rank_in_group
168
+ self.ep_size = self.ep_group.size()
169
+ self.n_routed_experts: int = config.n_routed_experts
170
+ self.n_shared_experts: int = config.n_shared_experts
171
+
172
+ self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe
173
+
174
+ if config.hidden_act != "silu":
175
+ raise ValueError(
176
+ f"Unsupported activation: {config.hidden_act}. "
177
+ "Only silu is supported for now."
178
+ )
179
+
180
+ self.router_dtype = _get_moe_router_dtype(config)
181
+ self.gate = GateLinear(
182
+ config.hidden_size,
183
+ config.n_routed_experts,
184
+ out_dtype=self.router_dtype,
185
+ prefix=f"{prefix}.gate",
186
+ )
187
+ if config.topk_method == "noaux_tc":
188
+ self.gate.e_score_correction_bias = nn.Parameter(
189
+ torch.empty(config.n_routed_experts, dtype=torch.float32)
190
+ )
191
+ else:
192
+ self.gate.e_score_correction_bias = None
193
+
194
+ # Load balancing settings.
195
+ eplb_config = parallel_config.eplb_config
196
+ self.enable_eplb = parallel_config.enable_eplb
197
+
198
+ self.n_redundant_experts = eplb_config.num_redundant_experts
199
+ self.n_logical_experts = self.n_routed_experts
200
+ self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts
201
+ self.n_local_physical_experts = self.n_physical_experts // self.ep_size
202
+
203
+ self.physical_expert_start = self.ep_rank * self.n_local_physical_experts
204
+ self.physical_expert_end = (
205
+ self.physical_expert_start + self.n_local_physical_experts
206
+ )
207
+
208
+ swiglu_limit = config.swiglu_limit
209
+ if config.n_shared_experts is None:
210
+ self.shared_experts = None
211
+ else:
212
+ intermediate_size = config.moe_intermediate_size * config.n_shared_experts
213
+
214
+ self.shared_experts = Glm5NextMLP(
215
+ hidden_size=config.hidden_size,
216
+ intermediate_size=intermediate_size,
217
+ hidden_act=config.hidden_act,
218
+ quant_config=quant_config,
219
+ is_sequence_parallel=self.is_sequence_parallel,
220
+ reduce_results=False,
221
+ prefix=f"{prefix}.shared_experts",
222
+ swiglu_limit=swiglu_limit,
223
+ )
224
+
225
+ self.experts = FusedMoEFactory(
226
+ shared_experts=self.shared_experts,
227
+ gate=self.gate,
228
+ num_experts=config.n_routed_experts,
229
+ top_k=config.num_experts_per_token,
230
+ hidden_size=config.hidden_size,
231
+ intermediate_size=config.moe_intermediate_size,
232
+ renormalize=config.moe_renormalize,
233
+ quant_config=quant_config,
234
+ use_grouped_topk=True,
235
+ num_expert_group=config.n_group,
236
+ topk_group=config.topk_group,
237
+ prefix=f"{prefix}.experts",
238
+ scoring_func=config.scoring_func,
239
+ routed_scaling_factor=self.routed_scaling_factor,
240
+ apply_routed_scale_to_output=apply_routed_scale_to_output,
241
+ e_score_correction_bias=self.gate.e_score_correction_bias,
242
+ enable_eplb=self.enable_eplb,
243
+ num_redundant_experts=self.n_redundant_experts,
244
+ is_sequence_parallel=self.is_sequence_parallel,
245
+ n_shared_experts=None,
246
+ router_logits_dtype=self.gate.out_dtype,
247
+ swiglu_limit=swiglu_limit,
248
+ )
249
+
250
+ def forward(
251
+ self,
252
+ hidden_states: torch.Tensor,
253
+ already_sequence_parallel: bool = False,
254
+ ) -> torch.Tensor:
255
+ num_tokens, hidden_dim = hidden_states.shape
256
+
257
+ # Chunk the hidden states so they aren't replicated across TP ranks.
258
+ # This avoids duplicate computation in self.experts.
259
+ if self.is_sequence_parallel and not already_sequence_parallel:
260
+ hidden_states = sequence_parallel_chunk(hidden_states)
261
+
262
+ # The router is always external (self.gate); main's MoERunner expects
263
+ # pre-computed router_logits, so compute them here unconditionally.
264
+ router_logits, _ = self.gate(hidden_states)
265
+ final_hidden_states = self.experts(
266
+ hidden_states=hidden_states, router_logits=router_logits
267
+ )
268
+
269
+ if self.is_sequence_parallel and not already_sequence_parallel:
270
+ final_hidden_states = tensor_model_parallel_all_gather(
271
+ final_hidden_states, 0
272
+ )
273
+ final_hidden_states = final_hidden_states[:num_tokens]
274
+
275
+ return final_hidden_states.view(num_tokens, hidden_dim)
276
+
277
+
278
+class Glm5NextDecoderLayer(nn.Module):
279
+ def __init__(
280
+ self,
281
+ vllm_config: VllmConfig,
282
+ config: Glm5NextConfig,
283
+ layer_idx: int,
284
+ prefix: str = "",
285
+ topk_indices_buffer: torch.Tensor | None = None,
286
+ is_mtp_layer: bool = False,
287
+ **kwargs,
288
+ ) -> None:
289
+ super().__init__()
290
+
291
+ cache_config = vllm_config.cache_config
292
+ quant_config = vllm_config.quant_config
293
+ parallel_config = vllm_config.parallel_config
294
+
295
+ self.hidden_size = config.hidden_size
296
+ self.layer_idx = layer_idx
297
+ self.is_moe = config.is_moe
298
+ self.num_hidden_layers = config.num_hidden_layers
299
+ self.rms_norm_eps = config.rms_norm_eps
300
+ self.num_experts = config.n_routed_experts
301
+ self.is_mtp_layer = is_mtp_layer
302
+ self.mhc = config.mhc
303
+ is_kda_layer = not is_mtp_layer and config.is_kda_layer(layer_idx)
304
+ self.layer_kind = "kda" if is_kda_layer else "mla"
305
+ self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe
306
+
307
+ if is_kda_layer:
308
+ self.self_attn = Glm5NextLinearAttention(
309
+ config=config,
310
+ vllm_config=vllm_config,
311
+ prefix=f"{prefix}.self_attn",
312
+ )
313
+ else:
314
+ # MLA layers require the latent head dims, which are guaranteed set
315
+ # on MLA configs; narrow away the `int | None`.
316
+ assert config.v_head_dim is not None
317
+ assert config.kv_lora_rank is not None
318
+ self.self_attn = Glm5NextMLAAttention(
319
+ vllm_config=vllm_config,
320
+ config=config,
321
+ hidden_size=self.hidden_size,
322
+ num_heads=config.num_attention_heads,
323
+ qk_nope_head_dim=config.qk_nope_head_dim,
324
+ qk_rope_head_dim=config.qk_rope_head_dim,
325
+ v_head_dim=config.v_head_dim,
326
+ q_lora_rank=config.q_lora_rank,
327
+ kv_lora_rank=config.kv_lora_rank,
328
+ max_position_embeddings=config.max_position_embeddings,
329
+ cache_config=cache_config,
330
+ quant_config=None, # MLA projections are BF16 in checkpoint
331
+ prefix=f"{prefix}.self_attn",
332
+ topk_indices_buffer=topk_indices_buffer,
333
+ skip_rope=config.mla_nope,
334
+ )
335
+
336
+ # MTP layers sit past the base model's hidden layers (layer_idx >=
337
+ # num_hidden_layers), so they're outside mlp_layer_types; default them
338
+ # to the last base layer's MLP type (sparse/MoE for these checkpoints).
339
+ mlp_layer_types = config.mlp_layer_types
340
+ mlp_type = (
341
+ mlp_layer_types[layer_idx]
342
+ if layer_idx < len(mlp_layer_types)
343
+ else (mlp_layer_types[-1] if mlp_layer_types else "sparse")
344
+ )
345
+ if self.is_moe and self.num_experts is not None and mlp_type == "sparse":
346
+ self.mlp = Glm5NextMoE(
347
+ config=config,
348
+ parallel_config=parallel_config,
349
+ quant_config=quant_config,
350
+ prefix=f"{prefix}.mlp",
351
+ )
352
+ else:
353
+ self.mlp = Glm5NextMLP(
354
+ hidden_size=self.hidden_size,
355
+ intermediate_size=config.intermediate_size,
356
+ hidden_act=config.hidden_act,
357
+ quant_config=quant_config,
358
+ prefix=f"{prefix}.mlp",
359
+ swiglu_limit=config.swiglu_limit,
360
+ )
361
+ self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
362
+ # Cached for the hot forward path (isinstance per layer per step).
363
+ self._mlp_is_moe = isinstance(self.mlp, Glm5NextMoE)
364
+ # In SP, the attention output projection leaves a partial sum; the
365
+ # decoder-layer reduce_scatter after attention completes it (DSv4 pattern).
366
+ # MTP layers use the non-mHC path which has no sp_reduce_scatter, so
367
+ # their o_proj must still reduce normally.
368
+ if self.is_sequence_parallel and not is_mtp_layer:
369
+ self.self_attn.o_proj.reduce_results = False
370
+ self.post_attention_layernorm = RMSNorm(
371
+ config.hidden_size, eps=config.rms_norm_eps
372
+ )
373
+
374
+ if self.mhc and not is_mtp_layer:
375
+ # mhc config
376
+ self.mhc_num_residual_streams = config.mhc_num_residual_streams
377
+ self.mhc_tau = config.mhc_tau
378
+ self.hc_eps = config.hc_eps
379
+ self.mhc_sinkhorn_iterations = config.mhc_sinkhorn_iterations
380
+ self.mhc_post_mult_value = config.mhc_post_mult_value
381
+
382
+ n = config.mhc_num_residual_streams
383
+ d_model = n * self.hidden_size
384
+ mix_hc = (2 + n) * n
385
+
386
+ self.n = n
387
+
388
+ # attn hc
389
+ self.hc_attn_fn = nn.Parameter(
390
+ torch.empty(mix_hc, d_model, dtype=torch.float32)
391
+ )
392
+ self.hc_attn_base = nn.Parameter(torch.empty(mix_hc, dtype=torch.float32))
393
+ self.hc_attn_scale = nn.Parameter(torch.empty(3, dtype=torch.float32))
394
+
395
+ # ffn hc
396
+ self.hc_ffn_fn = nn.Parameter(
397
+ torch.empty(mix_hc, d_model, dtype=torch.float32)
398
+ )
399
+ self.hc_ffn_base = nn.Parameter(torch.empty(mix_hc, dtype=torch.float32))
400
+ self.hc_ffn_scale = nn.Parameter(torch.empty(3, dtype=torch.float32))
401
+
402
+ self.mhc_pre_op = MHCPreOp()
403
+ self.mhc_post_op = MHCPostOp()
404
+ self.mhc_fused_post_pre_op = MHCFusedPostPreOp()
405
+
406
+ def forward(
407
+ self,
408
+ positions: torch.Tensor,
409
+ hidden_states: torch.Tensor,
410
+ residual: torch.Tensor | None = None,
411
+ post: torch.Tensor | None = None,
412
+ comb: torch.Tensor | None = None,
413
+ ) -> tuple[
414
+ torch.Tensor,
415
+ torch.Tensor | None,
416
+ torch.Tensor | None,
417
+ torch.Tensor | None,
418
+ ]:
419
+ # 70B or MTP layers: KDA + MoE without HC.
420
+ if not self.mhc or self.is_mtp_layer:
421
+ residual = hidden_states
422
+ hidden_states = self.input_layernorm(hidden_states)
423
+
424
+ attn_output = self.self_attn(
425
+ hidden_states=hidden_states,
426
+ positions=positions,
427
+ )
428
+ hidden_states, residual = self.post_attention_layernorm(
429
+ attn_output, residual=residual
430
+ )
431
+ hidden_states = self.mlp(hidden_states)
432
+ if self.is_mtp_layer:
433
+ # Return the unsummed pair: the MTP caller feeds it straight
434
+ # into shared_head's fused_add_rms_norm (one kernel instead of
435
+ # a separate residual-add + norm). The sum itself is unchanged
436
+ # (fp32-accumulated inside the fused kernel).
437
+ return hidden_states, residual, None, None
438
+ hidden_states = residual + hidden_states
439
+ return hidden_states, residual, None, None
440
+
441
+ # mHC start. `post`/`comb` carry the previous layer's deferred
442
+ # hc_post inputs (its ffn-pre outputs); when present, fuse that
443
+ # hc_post with this layer's attn hc_pre into one kernel (inter-layer
444
+ # fusion). Layer 0 has no incoming state -> standalone hc_pre.
445
+ x = hidden_states
446
+ if post is None:
447
+ if self.layer_idx == 0:
448
+ x = hc_expand(x, self.n)
449
+ residual = x
450
+ post, comb, x = self.hc_pre(
451
+ x,
452
+ self.hc_attn_fn,
453
+ self.hc_attn_scale,
454
+ self.hc_attn_base,
455
+ norm_weight=self.input_layernorm.weight.data,
456
+ norm_eps=self.input_layernorm.variance_epsilon,
457
+ )
458
+ else:
459
+ residual, post, comb, x = self.hc_fused_post_pre(
460
+ x,
461
+ residual,
462
+ post,
463
+ comb,
464
+ self.hc_attn_fn,
465
+ self.hc_attn_scale,
466
+ self.hc_attn_base,
467
+ norm_weight=self.input_layernorm.weight.data,
468
+ norm_eps=self.input_layernorm.variance_epsilon,
469
+ )
470
+
471
+ # Attention needs the full token sequence; mHC above ran on the SP
472
+ # shard. Gather for attention, scatter back afterward (DSv4 pattern).
473
+ if self.is_sequence_parallel:
474
+ x = sp_all_gather(x)[: positions.shape[0]]
475
+
476
+ x = self.self_attn(
477
+ hidden_states=x,
478
+ positions=positions,
479
+ )
480
+
481
+ if self.is_sequence_parallel:
482
+ x = sp_reduce_scatter(x)
483
+
484
+ # Fuse post-attn hc_post + pre-FFN hc_pre (+ RMSNorm) into one kernel.
485
+ residual, post, comb, x = self.hc_fused_post_pre(
486
+ x,
487
+ residual,
488
+ post,
489
+ comb,
490
+ self.hc_ffn_fn,
491
+ self.hc_ffn_scale,
492
+ self.hc_ffn_base,
493
+ norm_weight=self.post_attention_layernorm.weight.data,
494
+ norm_eps=self.post_attention_layernorm.variance_epsilon,
495
+ )
496
+
497
+ # Fully Connected
498
+ if self._mlp_is_moe:
499
+ x = self.mlp(x, already_sequence_parallel=self.is_sequence_parallel)
500
+ else:
501
+ x = self.mlp(x)
502
+
503
+ # mHC end. The last mHC layer materializes its final hc_post (nothing
504
+ # to fuse with) then contracts; every other layer defers its hc_post to
505
+ # the next layer's fused pre, returning the state.
506
+ if self.layer_idx == self.num_hidden_layers - 1:
507
+ x = self.hc_post(x, residual, post, comb)
508
+ x = hc_contract(x, self.n)
509
+ return x, None, None, None
510
+
511
+ return x, residual, post, comb
512
+
513
+ def hc_pre(
514
+ self,
515
+ x: torch.Tensor,
516
+ hc_fn: torch.Tensor,
517
+ hc_scale: torch.Tensor,
518
+ hc_base: torch.Tensor,
519
+ norm_weight: torch.Tensor | None = None,
520
+ norm_eps: float = 0.0,
521
+ ):
522
+ post_mix, res_mix, layer_input = self.mhc_pre_op(
523
+ residual=x,
524
+ fn=hc_fn,
525
+ hc_scale=hc_scale,
526
+ hc_base=hc_base,
527
+ rms_eps=self.rms_norm_eps,
528
+ hc_pre_eps=self.hc_eps,
529
+ hc_sinkhorn_eps=self.hc_eps,
530
+ hc_post_mult_value=self.mhc_post_mult_value,
531
+ sinkhorn_repeat=self.mhc_sinkhorn_iterations,
532
+ norm_weight=norm_weight,
533
+ norm_eps=norm_eps,
534
+ )
535
+ return post_mix, res_mix, layer_input
536
+
537
+ def hc_post(
538
+ self,
539
+ x: torch.Tensor,
540
+ residual: torch.Tensor,
541
+ post: torch.Tensor,
542
+ comb: torch.Tensor,
543
+ ):
544
+ return self.mhc_post_op(x, residual, post, comb)
545
+
546
+ def hc_fused_post_pre(
547
+ self,
548
+ x: torch.Tensor,
549
+ residual: torch.Tensor,
550
+ post: torch.Tensor,
551
+ comb: torch.Tensor,
552
+ hc_fn: torch.Tensor,
553
+ hc_scale: torch.Tensor,
554
+ hc_base: torch.Tensor,
555
+ norm_weight: torch.Tensor | None = None,
556
+ norm_eps: float = 0.0,
557
+ ):
558
+ return self.mhc_fused_post_pre_op(
559
+ x=x,
560
+ residual=residual,
561
+ post_layer_mix=post,
562
+ comb_res_mix=comb,
563
+ fn=hc_fn,
564
+ hc_scale=hc_scale,
565
+ hc_base=hc_base,
566
+ rms_eps=self.rms_norm_eps,
567
+ hc_pre_eps=self.hc_eps,
568
+ hc_sinkhorn_eps=self.hc_eps,
569
+ hc_post_mult_value=self.mhc_post_mult_value,
570
+ sinkhorn_repeat=self.mhc_sinkhorn_iterations,
571
+ n_splits=1,
572
+ tile_n=1,
573
+ norm_weight=norm_weight,
574
+ norm_eps=norm_eps,
575
+ )
576
+
577
+
578
+class Glm5NextModel(nn.Module):
579
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
580
+ super().__init__()
581
+
582
+ config = vllm_config.model_config.hf_config
583
+ self.config = config
584
+
585
+ self.vocab_size = config.vocab_size
586
+ self.device = current_platform.device_type
587
+
588
+ self.is_v32 = config.index_topk is not None
589
+ if self.is_v32:
590
+ topk_tokens = config.index_topk
591
+ assert topk_tokens is not None
592
+ # Reserve room for the incomplete pool tail.
593
+ kpool = config.index_kpool
594
+ assert kpool is not None
595
+ buffer_width = topk_tokens + (kpool - 1 if kpool > 1 else 0)
596
+ # Sparse MLA tiles top-k in 128 columns; padded slots remain masked.
597
+ sparse_topk_block_n = 128
598
+ buffer_width = (
599
+ (buffer_width + sparse_topk_block_n - 1) // sparse_topk_block_n
600
+ ) * sparse_topk_block_n
601
+ topk_indices_buffer = torch.empty(
602
+ vllm_config.scheduler_config.max_num_batched_tokens,
603
+ buffer_width,
604
+ dtype=torch.int32,
605
+ device=self.device,
606
+ )
607
+ else:
608
+ # Full-MLA config (no kpool sparse indexer): no topk buffer.
609
+ topk_indices_buffer = None
610
+
611
+ if get_pp_group().is_first_rank:
612
+ self.embed_tokens = VocabParallelEmbedding(
613
+ config.vocab_size,
614
+ config.hidden_size,
615
+ prefix=f"{prefix}.embed_tokens",
616
+ )
617
+ else:
618
+ self.embed_tokens = PPMissingLayer()
619
+
620
+ def get_layer(prefix: str):
621
+ layer_idx = int(prefix.rsplit(".", 1)[1])
622
+ return Glm5NextDecoderLayer(
623
+ vllm_config=vllm_config,
624
+ config=config,
625
+ layer_idx=layer_idx,
626
+ prefix=prefix,
627
+ topk_indices_buffer=topk_indices_buffer,
628
+ )
629
+
630
+ self.start_layer, self.end_layer, self.layers = make_layers(
631
+ config.num_hidden_layers,
632
+ get_layer,
633
+ prefix=f"{prefix}.layers",
634
+ )
635
+ # The active slice is fixed after construction; cache it so forward
636
+ # doesn't rebuild the slice (a fresh list) every step.
637
+ self._active_layers = self.layers[self.start_layer : self.end_layer]
638
+
639
+ if get_pp_group().is_last_rank:
640
+ self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
641
+ else:
642
+ self.norm = PPMissingLayer()
643
+
644
+ self.is_sequence_parallel = (
645
+ vllm_config.parallel_config.use_sequence_parallel_moe
646
+ )
647
+
648
+ world_size = get_tensor_model_parallel_world_size()
649
+ assert config.num_attention_heads % world_size == 0, (
650
+ "num_attention_heads must be divisible by world_size"
651
+ )
652
+
653
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
654
+ return self.embed_tokens(input_ids)
655
+
656
+ def forward(
657
+ self,
658
+ input_ids: torch.Tensor | None,
659
+ positions: torch.Tensor,
660
+ intermediate_tensors: IntermediateTensors | None,
661
+ inputs_embeds: torch.Tensor | None = None,
662
+ **kwargs,
663
+ ) -> torch.Tensor:
664
+ if get_pp_group().is_first_rank:
665
+ if inputs_embeds is not None:
666
+ hidden_states = inputs_embeds
667
+ else:
668
+ hidden_states = self.embed_input_ids(input_ids)
669
+ residual = None
670
+ post = None
671
+ comb = None
672
+ else:
673
+ assert intermediate_tensors is not None
674
+ hidden_states = intermediate_tensors["hidden_states"]
675
+ residual = intermediate_tensors["residual"]
676
+ # post/comb (deferred mHC hc_post state) are not propagated across
677
+ # PP ranks; the receiving rank's first mHC layer uses standalone pre.
678
+ post = None
679
+ comb = None
680
+
681
+ full_num_tokens = positions.shape[0]
682
+ if self.is_sequence_parallel:
683
+ hidden_states = sp_shard(hidden_states)
684
+
685
+ for layer in self._active_layers:
686
+ hidden_states, residual, post, comb = layer(
687
+ positions, hidden_states, residual, post, comb
688
+ )
689
+
690
+ if not get_pp_group().is_last_rank:
691
+ # PP is gated off for GLM-5.3-Flash (no make_empty_intermediate_tensors),
692
+ # so this branch is not exercised. post/comb are the deferred
693
+ # hc_post state of this rank's last mHC layer; a future PP path
694
+ # would need to propagate them, but for now they are dropped (the
695
+ # receiving rank's first layer would fall back to standalone pre).
696
+ return IntermediateTensors(
697
+ {"hidden_states": hidden_states, "residual": residual}
698
+ )
699
+
700
+ if self.is_sequence_parallel:
701
+ hidden_states = sp_all_gather(hidden_states)[:full_num_tokens]
702
+
703
+ hidden_states = self.norm(hidden_states)
704
+ return hidden_states
705
+
706
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
707
+ stacked_params_mapping = [
708
+ # (param_name, shard_name, shard_id)
709
+ (".gate_up_proj", ".gate_proj", 0),
710
+ (".gate_up_proj", ".up_proj", 1),
711
+ # MLA: fuse q_a_proj and kv_a_proj_with_mqa
712
+ (".fused_qkv_a_proj", ".q_a_proj", 0),
713
+ (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1),
714
+ # Indexer: fuse wk and weights_proj
715
+ (".wk_weights_proj", ".wk", 0),
716
+ (".wk_weights_proj", ".weights_proj", 1),
717
+ # KDA: merge q, k, v, b, f_a, g_a projections into one GEMM
718
+ (".in_proj_qkvbfg_a", ".q_proj", 0),
719
+ (".in_proj_qkvbfg_a", ".k_proj", 1),
720
+ (".in_proj_qkvbfg_a", ".v_proj", 2),
721
+ (".in_proj_qkvbfg_a", ".b_proj", 3),
722
+ (".in_proj_qkvbfg_a", ".f_a_proj", 4),
723
+ (".in_proj_qkvbfg_a", ".g_a_proj", 5),
724
+ ]
725
+ if self.config.is_moe:
726
+ # Params for weights, fp8 weight scales, fp8 activation scales
727
+ # (param_name, weight_name, expert_id, shard_id)
728
+ expert_params_mapping = fused_moe_make_expert_params_mapping(
729
+ self,
730
+ ckpt_gate_proj_name="gate_proj",
731
+ ckpt_down_proj_name="down_proj",
732
+ ckpt_up_proj_name="up_proj",
733
+ num_experts=self.config.n_routed_experts,
734
+ )
735
+ else:
736
+ expert_params_mapping = []
737
+ params_dict = dict(self.named_parameters())
738
+ loaded_params: set[str] = set()
739
+
740
+ # GLM-5.3-Flash NoPE checkpoints omit the RoPE rows from
741
+ # ``kv_a_proj_with_mqa``; pad them with zeros for the model shape.
742
+ kv_a_pad_size = 0
743
+ if self.config.mla_nope and self.config.qk_rope_head_dim > 0:
744
+ kv_a_pad_size = self.config.qk_rope_head_dim
745
+
746
+ _pending_wk_fp8: dict = {}
747
+
748
+ for args in weights:
749
+ name, loaded_weight = args[:2]
750
+ kwargs: dict = args[2] if len(args) > 2 else {}
751
+ if "rotary_emb.inv_freq" in name:
752
+ continue
753
+
754
+ spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
755
+ if spec_layer is not None:
756
+ continue # skip spec decode layers for main model
757
+ if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
758
+ # Models trained using ColossalAI may include these tensors in
759
+ # the checkpoint. Skip them.
760
+ continue
761
+
762
+ # Handle FP8 indexer WK: dequantize to BF16 for fusion with
763
+ # weights_proj into wk_weights_proj.
764
+ if _try_load_fp8_indexer_wk(
765
+ name,
766
+ loaded_weight,
767
+ _pending_wk_fp8,
768
+ params_dict,
769
+ loaded_params,
770
+ ):
771
+ continue
772
+
773
+ # FP8 checkpoint: dequantize BF16-kept MLA projections
774
+ # (q_a_proj / kv_a_proj_with_mqa / o_proj) to BF16.
775
+ if _try_load_fp8_attn_proj(
776
+ name,
777
+ loaded_weight,
778
+ _pending_wk_fp8,
779
+ params_dict,
780
+ loaded_params,
781
+ kv_a_pad_size,
782
+ ):
783
+ continue
784
+
785
+ # Pad kv_a_proj_with_mqa for NoPE models
786
+ if kv_a_pad_size > 0 and ".kv_a_proj_with_mqa." in name:
787
+ pad = torch.zeros(
788
+ kv_a_pad_size,
789
+ *loaded_weight.shape[1:],
790
+ dtype=loaded_weight.dtype,
791
+ device=loaded_weight.device,
792
+ )
793
+ loaded_weight = torch.cat([loaded_weight, pad], dim=0)
794
+
795
+ for param_name, weight_name, shard_id in stacked_params_mapping:
796
+ if weight_name not in name:
797
+ continue
798
+ # We have mlp.experts[0].gate_proj in the checkpoint.
799
+ # Since we handle the experts below in expert_params_mapping,
800
+ # we need to skip here BEFORE we update the name, otherwise
801
+ # name will be updated to mlp.experts[0].gate_up_proj, which
802
+ # will then be updated below in expert_params_mapping
803
+ # for mlp.experts[0].gate_gate_up_proj, which breaks load.
804
+ if ("mlp.experts." in name) and name not in params_dict:
805
+ continue
806
+ name_mapped = name.replace(weight_name, param_name)
807
+ # QKV fusion: skip if fused module doesn't exist in model
808
+ if param_name == ".fused_qkv_a_proj" and name_mapped not in params_dict:
809
+ continue
810
+ name = name_mapped
811
+ # Skip loading extra bias for GPTQ models.
812
+ if name.endswith(".bias") and name not in params_dict:
813
+ continue
814
+ if is_pp_missing_parameter(name, self):
815
+ continue
816
+ param = params_dict[name]
817
+ weight_loader = param.weight_loader
818
+ weight_loader(param, loaded_weight, shard_id)
819
+ break
820
+ else:
821
+ for idx, (
822
+ param_name,
823
+ weight_name,
824
+ expert_id,
825
+ expert_shard_id,
826
+ ) in enumerate(expert_params_mapping):
827
+ if weight_name not in name:
828
+ continue
829
+ name = name.replace(weight_name, param_name)
830
+ if is_pp_missing_parameter(name, self):
831
+ continue
832
+ param = params_dict[name]
833
+ weight_loader = param.weight_loader
834
+ weight_loader(
835
+ param,
836
+ loaded_weight,
837
+ name,
838
+ expert_id=expert_id,
839
+ shard_id=expert_shard_id,
840
+ )
841
+ break
842
+ else:
843
+ # Skip loading extra bias for GPTQ models.
844
+ if (
845
+ name.endswith(".bias")
846
+ and name not in params_dict
847
+ and not self.config.is_linear_attn
848
+ ): # noqa: E501
849
+ continue
850
+ # Remapping the name of FP8 kv-scale.
851
+ remapped_name = maybe_remap_kv_scale_name(name, params_dict)
852
+ if remapped_name is None:
853
+ continue
854
+ name = remapped_name
855
+ if is_pp_missing_parameter(name, self):
856
+ continue
857
+
858
+ param = params_dict[name]
859
+ weight_loader = getattr(
860
+ param, "weight_loader", default_weight_loader
861
+ )
862
+ weight_loader(param, loaded_weight, **kwargs)
863
+ loaded_params.add(name)
864
+ return loaded_params
865
+
866
+
867
+class Glm5NextForCausalLM(
868
+ nn.Module, HasInnerState, SupportsPP, MixtureOfExperts, IsHybrid
869
+):
870
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
871
+ super().__init__()
872
+ self.model_config = vllm_config.model_config
873
+ self.vllm_config = vllm_config
874
+ self.config = self.model_config.hf_config
875
+ quant_config = vllm_config.quant_config
876
+ self.quant_config = quant_config
877
+ self.model = Glm5NextModel(
878
+ vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
879
+ )
880
+ if get_pp_group().is_last_rank:
881
+ self.lm_head = ParallelLMHead(
882
+ self.config.vocab_size,
883
+ self.config.hidden_size,
884
+ quant_config=quant_config,
885
+ prefix=maybe_prefix(prefix, "lm_head"),
886
+ )
887
+ else:
888
+ self.lm_head = PPMissingLayer()
889
+ self.logits_processor = LogitsProcessor(
890
+ self.config.vocab_size, scale=self.config.logit_scale
891
+ )
892
+
893
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
894
+ return self.model.embed_input_ids(input_ids)
895
+
896
+ def forward(
897
+ self,
898
+ input_ids: torch.Tensor | None,
899
+ positions: torch.Tensor,
900
+ intermediate_tensors: IntermediateTensors | None = None,
901
+ inputs_embeds: torch.Tensor | None = None,
902
+ **kwargs,
903
+ ) -> torch.Tensor | IntermediateTensors:
904
+ hidden_states = self.model(
905
+ input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs
906
+ )
907
+ return hidden_states
908
+
909
+ @classmethod
910
+ def get_mamba_state_dtype_from_config(
911
+ cls,
912
+ vllm_config: "VllmConfig",
913
+ ) -> tuple[torch.dtype, torch.dtype]:
914
+ return MambaStateDtypeCalculator.kda_state_dtype(
915
+ vllm_config.model_config.dtype, vllm_config.cache_config.mamba_cache_dtype
916
+ )
917
+
918
+ @classmethod
919
+ def get_mamba_state_shape_from_config(
920
+ cls, vllm_config: "VllmConfig"
921
+ ) -> tuple[tuple[int, int], tuple[int, int, int]]:
922
+ parallel_config = vllm_config.parallel_config
923
+ hf_config = vllm_config.model_config.hf_config
924
+ tp_size = parallel_config.tensor_parallel_size
925
+ num_spec = (
926
+ vllm_config.speculative_config.num_speculative_tokens
927
+ if vllm_config.speculative_config
928
+ else 0
929
+ )
930
+ return MambaStateShapeCalculator.kda_state_shape(
931
+ tp_size,
932
+ hf_config.linear_num_heads,
933
+ hf_config.linear_head_dim,
934
+ conv_kernel_size=hf_config.linear_conv_kernel_dim,
935
+ num_spec=num_spec,
936
+ )
937
+
938
+ @classmethod
939
+ def get_mamba_state_copy_func(
940
+ cls,
941
+ ) -> tuple[
942
+ MambaStateCopyFunc, MambaStateCopyFunc, MambaStateCopyFunc, MambaStateCopyFunc
943
+ ]:
944
+ return MambaStateCopyFuncCalculator.kda_state_copy_func()
945
+
946
+ def compute_logits(
947
+ self,
948
+ hidden_states: torch.Tensor,
949
+ ) -> torch.Tensor | None:
950
+ logits = self.logits_processor(self.lm_head, hidden_states)
951
+ return logits
952
+
953
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
954
+ loader = AutoWeightsLoader(self)
955
+ return loader.load_weights(weights)
956
+
957
+
958
+@MULTIMODAL_REGISTRY.register_processor(
959
+ Glm5NextMultiModalProcessor,
960
+ info=Glm5NextProcessingInfo,
961
+ dummy_inputs=Glm4vDummyInputsBuilder,
962
+)
963
+class Glm5NextForConditionalGeneration(
964
+ Glm4vForConditionalGeneration, HasInnerState, IsHybrid
965
+):
966
+ # The text model (KDA + dense-MLA + MoE) is a hybrid mamba model. The
967
+ # multimodal wrapper must declare the same interfaces so vLLM treats it as
968
+ # hybrid (auto-aligns mamba/attention block sizes, sizes the mamba state
969
+ # cache); the mamba-state classmethods delegate to the text model.
970
+ has_inner_state: ClassVar[Literal[True]] = True
971
+ is_hybrid: ClassVar[Literal[True]] = True
972
+
973
+ # NOTE: weight-prefix mapping is inherited from Glm4vForConditionalGeneration
974
+ # (``model.visual.`` -> ``visual.``, ``model.language_model.`` ->
975
+ # ``language_model.model.``, ``lm_head.`` -> ``language_model.lm_head.``),
976
+ # matching the GLM-OCR / GLM-4V serialization convention. If the real
977
+ # checkpoint's safetensors keys differ (e.g. ``language_model.model.`` with
978
+ # no outer ``model.``), override ``hf_to_vllm_mapper`` accordingly.
979
+
980
+ @classmethod
981
+ def get_mamba_state_dtype_from_config(cls, vllm_config: VllmConfig):
982
+ from .model import Glm5NextForCausalLM
983
+
984
+ return Glm5NextForCausalLM.get_mamba_state_dtype_from_config(vllm_config)
985
+
986
+ @classmethod
987
+ def get_mamba_state_shape_from_config(cls, vllm_config: VllmConfig):
988
+ from .model import Glm5NextForCausalLM
989
+
990
+ return Glm5NextForCausalLM.get_mamba_state_shape_from_config(vllm_config)
991
+
992
+ @classmethod
993
+ def get_mamba_state_copy_func(cls):
994
+ from .model import Glm5NextForCausalLM
995
+
996
+ return Glm5NextForCausalLM.get_mamba_state_copy_func()
997
+
998
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
999
+ super(Glm4vForConditionalGeneration, self).__init__()
1000
+ config = vllm_config.model_config.hf_config
1001
+ multimodal_config = vllm_config.model_config.multimodal_config
1002
+ assert multimodal_config is not None
1003
+
1004
+ self.config = config
1005
+ self.model_config = vllm_config.model_config
1006
+ self.multimodal_config = multimodal_config
1007
+ self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"
1008
+ self.is_multimodal_pruning_enabled = (
1009
+ multimodal_config.is_multimodal_pruning_enabled()
1010
+ )
1011
+
1012
+ with self._mark_tower_model(vllm_config, {"image", "video"}):
1013
+ self.visual = Glm5NextVisionTransformer(
1014
+ config.text_config,
1015
+ config.vision_config,
1016
+ # Read eps from the VISION sub-config, not the top-level
1017
+ # `config.rms_norm_eps`: Glm5NextConfig.__getattribute__ mirrors
1018
+ # the latter onto text_config (1e-5), silently ignoring the
1019
+ # vision tower's own (1e-6) rms_norm_eps.
1020
+ norm_eps=config.vision_config.rms_norm_eps,
1021
+ # Vision tower ships BF16 weights in this fp8 checkpoint (no
1022
+ # weight_scale_inv for visual.*), so it must NOT inherit the
1023
+ # global fp8 quant_config -- doing so incorrectly quantizes
1024
+ # the tower
1025
+ # and yields NaN image features. Mirrors the MLA/KDA proj
1026
+ # pattern (quant_config=None for BF16 submodules).
1027
+ quant_config=None,
1028
+ prefix=maybe_prefix(prefix, "visual"),
1029
+ )
1030
+
1031
+ with self._mark_language_model(vllm_config):
1032
+ self.language_model = init_vllm_registered_model(
1033
+ vllm_config=vllm_config,
1034
+ hf_config=config.text_config,
1035
+ prefix=maybe_prefix(prefix, "language_model"),
1036
+ architectures=["Glm5NextForCausalLM"],
1037
+ )
1038
+
1039
+ # Glm5NextForCausalLM does not implement make_empty_intermediate_tensors,
1040
+ # so pipeline parallelism is gated off (consistent with the text-only
1041
+ # model) and we intentionally do not alias it here.
1042
+
1043
+ def get_encoder_cudagraph_config(self):
1044
+ # This vision tower does not produce the absolute position embedding
1045
+ # buffer used by GLM4V.
1046
+ config = super().get_encoder_cudagraph_config()
1047
+ config.buffer_keys = [k for k in config.buffer_keys if k != "pos_embeds"]
1048
+ return config
1049
+
1050
+
1051
+def get_spec_layer_idx_from_weight_name(
1052
+ config: Glm5NextConfig, weight_name: str
1053
+) -> int | None:
1054
+ if hasattr(config, "num_nextn_predict_layers") and (
1055
+ config.num_nextn_predict_layers > 0
1056
+ ):
1057
+ layer_idx = config.num_hidden_layers
1058
+ for i in range(config.num_nextn_predict_layers):
1059
+ if weight_name.startswith(
1060
+ f"model.layers.{layer_idx + i}."
1061
+ ) or weight_name.startswith(f"layers.{layer_idx + i}."):
1062
+ return layer_idx + i
1063
+ return None
1064
+
1065
+
1066
+def _try_load_fp8_indexer_wk(name, tensor, buf, params_dict, loaded_params):
1067
+ if "indexer.wk." not in name or "wk_weights" in name:
1068
+ return False
1069
+ is_weight = name.endswith(".weight") and tensor.dtype == torch.float8_e4m3fn
1070
+ is_scale = "weight_scale_inv" in name
1071
+ if not is_weight and not is_scale:
1072
+ return False
1073
+ layer_prefix = name.rsplit(".wk.", 1)[0]
1074
+ entry = buf.setdefault(layer_prefix, {})
1075
+ entry["weight" if is_weight else "scale"] = tensor
1076
+ if "weight" not in entry or "scale" not in entry:
1077
+ return True
1078
+
1079
+ weight_fp8, scale_inv = entry["weight"], entry["scale"]
1080
+ del buf[layer_prefix]
1081
+ block_size = weight_fp8.shape[1] // scale_inv.shape[1]
1082
+ weight_bf16 = scaled_dequantize(
1083
+ weight_fp8,
1084
+ scale_inv,
1085
+ group_shape=GroupShape(block_size, block_size),
1086
+ out_dtype=torch.bfloat16,
1087
+ )
1088
+
1089
+ fused_name = f"{layer_prefix}.wk_weights_proj.weight"
1090
+ param = params_dict[fused_name]
1091
+ param.weight_loader(param, weight_bf16, 0)
1092
+ loaded_params.add(fused_name)
1093
+ return True
1094
+
1095
+
1096
+def _dequant_fp8_block(
1097
+ weight_fp8: torch.Tensor,
1098
+ scale_inv: torch.Tensor,
1099
+ block_size: int = 128,
1100
+) -> torch.Tensor:
1101
+ """Dequantize a block-FP8 (e4m3) weight with per-block scale to BF16.
1102
+
1103
+ Unlike ``scaled_dequantize`` this tolerates a non-divisible (partial last
1104
+ block) shape by zero-padding to a multiple of ``block_size`` before the
1105
+ scale broadcast and trimming back afterwards (e.g. kv_a_proj_with_mqa is
1106
+ 576 rows = 4*128 + 64).
1107
+ """
1108
+ out_dim, in_dim = weight_fp8.shape
1109
+ pad_out = (-out_dim) % block_size
1110
+ pad_in = (-in_dim) % block_size
1111
+ w = weight_fp8
1112
+ if pad_out or pad_in:
1113
+ w = torch.nn.functional.pad(w, (0, pad_in, 0, pad_out))
1114
+ # scale_inv is (ceil(out/block), ceil(in/block)); broadcast to (out, in).
1115
+ s = scale_inv.to(torch.float32)
1116
+ s_full = s.repeat_interleave(block_size, dim=0).repeat_interleave(block_size, dim=1)
1117
+ out = (w.to(torch.float32) * s_full).to(torch.bfloat16)
1118
+ return out[:out_dim, :in_dim].contiguous()
1119
+
1120
+
1121
+# FP8 checkpoint projections that the MODEL keeps in BF16, so the block-FP8
1122
+# (weight + weight_scale_inv) must be dequantized to BF16 on load.
1123
+# Maps checkpoint proj-suffix -> (buffer key, model target base, fused shard id
1124
+# or None for a direct projection, whether NoPE rope-padding applies).
1125
+_FP8_ATTN_PROJS = {
1126
+ ".q_a_proj.": ("q_a", "fused_qkv_a_proj", 0, False),
1127
+ ".kv_a_proj_with_mqa.": ("kv_a", "fused_qkv_a_proj", 1, True),
1128
+ ".q_b_proj.": ("q_b", "q_b_proj", None, False),
1129
+ ".o_proj.": ("o_proj", "o_proj", None, False),
1130
+}
1131
+
1132
+
1133
+def _try_load_fp8_attn_proj(
1134
+ name,
1135
+ tensor,
1136
+ buf,
1137
+ params_dict,
1138
+ loaded_params,
1139
+ kv_a_pad_size: int,
1140
+) -> bool:
1141
+ """Dequantize FP8 q_a_proj / kv_a_proj_with_mqa / o_proj to BF16 on load.
1142
+
1143
+ The FP8 checkpoint stores these as block-FP8 (weight + weight_scale_inv),
1144
+ but the model holds them in BF16 (``fused_qkv_a_proj`` is always BF16 via
1145
+ DeepSeekV2FusedQkvAProjLinear; ``o_proj`` is excluded by
1146
+ modules_to_not_convert). When the model target is BF16 (no
1147
+ ``weight_scale_inv`` param) we dequantize; otherwise we return False so the
1148
+ normal stacked/direct path loads the FP8 tensor as-is.
1149
+ """
1150
+ matched = None
1151
+ for suffix, info in _FP8_ATTN_PROJS.items():
1152
+ if suffix in name:
1153
+ matched = (suffix, info)
1154
+ break
1155
+ if matched is None:
1156
+ return False
1157
+ suffix, (key, target_base, shard_id, is_kva) = matched
1158
+ is_weight = name.endswith(".weight") and tensor.dtype == torch.float8_e4m3fn
1159
+ is_scale = "weight_scale_inv" in name
1160
+ if not is_weight and not is_scale:
1161
+ return False
1162
+
1163
+ layer_prefix = name.rsplit(suffix, 1)[0]
1164
+ target_w = f"{layer_prefix}.{target_base}.weight"
1165
+ target_s = f"{layer_prefix}.{target_base}.weight_scale_inv"
1166
+ # If the model actually kept this projection in FP8, let the normal path
1167
+ # handle it (it has a weight_scale_inv param).
1168
+ if target_s in params_dict:
1169
+ return False
1170
+
1171
+ entry = buf.setdefault(layer_prefix, {}).setdefault(key, {})
1172
+ entry["weight" if is_weight else "scale"] = tensor
1173
+ if "weight" not in entry or "scale" not in entry:
1174
+ return True
1175
+
1176
+ weight_fp8, scale_inv = entry["weight"], entry["scale"]
1177
+ buf[layer_prefix].pop(key, None)
1178
+ block_size = weight_fp8.shape[1] // scale_inv.shape[1]
1179
+ weight_bf16 = _dequant_fp8_block(weight_fp8, scale_inv, block_size)
1180
+ # NoPE: pad kv_a rope portion (kv_lora_rank -> kv_lora_rank + qk_rope_head_dim).
1181
+ if is_kva and kv_a_pad_size > 0:
1182
+ pad = torch.zeros(
1183
+ kv_a_pad_size,
1184
+ weight_bf16.shape[1],
1185
+ dtype=weight_bf16.dtype,
1186
+ device=weight_bf16.device,
1187
+ )
1188
+ weight_bf16 = torch.cat([weight_bf16, pad], dim=0)
1189
+
1190
+ param = params_dict[target_w]
1191
+ if shard_id is None:
1192
+ param.weight_loader(param, weight_bf16)
1193
+ else:
1194
+ param.weight_loader(param, weight_bf16, shard_id)
1195
+ loaded_params.add(target_w)
1196
+ return True
vllm/models/glm5next/nvidia/mtp.py
new
+429
@@ -0,0 +1,429 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import typing
4
+from collections.abc import Callable, Iterable
5
+
6
+import torch
7
+import torch.nn as nn
8
+
9
+from vllm.config import VllmConfig
10
+from vllm.model_executor.layers.fused_moe import (
11
+ fused_moe_make_expert_params_mapping,
12
+)
13
+from vllm.model_executor.layers.layernorm import RMSNorm
14
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
15
+from vllm.model_executor.layers.vocab_parallel_embedding import (
16
+ VocabParallelEmbedding,
17
+)
18
+from vllm.model_executor.model_loader.weight_utils import (
19
+ default_weight_loader,
20
+ maybe_remap_kv_scale_name,
21
+)
22
+from vllm.model_executor.models.deepseek_mtp import SharedHead
23
+from vllm.model_executor.models.deepseek_v2 import DeepseekV2MixtureOfExperts
24
+from vllm.model_executor.models.utils import maybe_prefix
25
+from vllm.platforms import current_platform
26
+from vllm.sequence import IntermediateTensors
27
+
28
+from .model import (
29
+ Glm5NextDecoderLayer,
30
+ Glm5NextMLAAttention,
31
+ Glm5NextMoE,
32
+ _try_load_fp8_attn_proj,
33
+ _try_load_fp8_indexer_wk,
34
+ get_spec_layer_idx_from_weight_name,
35
+)
36
+from .ops.fused_eh_norm import fused_eh_norm
37
+
38
+
39
+class Glm5NextMultiTokenPredictorLayer(nn.Module):
40
+ def __init__(self, vllm_config: VllmConfig, prefix: str) -> None:
41
+ super().__init__()
42
+ assert vllm_config.speculative_config is not None
43
+ config = vllm_config.speculative_config.draft_model_config.hf_config
44
+ self.config = config
45
+ quant_config = vllm_config.quant_config
46
+
47
+ self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
48
+ self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
49
+ self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False)
50
+
51
+ # Reserve room for the incomplete pool tail and align the sparse MLA
52
+ # buffer width to BLOCK_N=128.
53
+ topk_tokens = config.index_topk
54
+ assert topk_tokens is not None
55
+ kpool = config.index_kpool
56
+ assert kpool is not None
57
+ buffer_width = topk_tokens + (kpool - 1 if kpool > 1 else 0)
58
+ sparse_topk_block_n = 128
59
+ buffer_width = (
60
+ (buffer_width + sparse_topk_block_n - 1) // sparse_topk_block_n
61
+ ) * sparse_topk_block_n
62
+ topk_indices_buffer = torch.empty(
63
+ vllm_config.scheduler_config.max_num_batched_tokens,
64
+ buffer_width,
65
+ dtype=torch.int32,
66
+ device=current_platform.device_type,
67
+ )
68
+ self.shared_head = SharedHead(
69
+ config=config, prefix=prefix, quant_config=quant_config
70
+ )
71
+ # MTP layers sit past the base model's hidden layers; parse the index
72
+ # from the prefix (e.g. "...layers.32") so the decoder builds an MLA
73
+ # (DSA) layer rather than KDA for the MTP path.
74
+ layer_idx = int(prefix.rsplit(".", 1)[-1])
75
+ self.mtp_block = Glm5NextDecoderLayer(
76
+ vllm_config=vllm_config,
77
+ config=config,
78
+ layer_idx=layer_idx,
79
+ prefix=prefix,
80
+ topk_indices_buffer=topk_indices_buffer,
81
+ is_mtp_layer=True,
82
+ )
83
+
84
+ def forward(
85
+ self,
86
+ input_ids: torch.Tensor,
87
+ positions: torch.Tensor,
88
+ previous_hidden_states: torch.Tensor,
89
+ inputs_embeds: torch.Tensor | None = None,
90
+ spec_step_index: int = 0,
91
+ ) -> torch.Tensor:
92
+ assert inputs_embeds is not None
93
+ # Fused: zero pos-0 embeds + enorm(embeds) + hnorm(prev) + cat -> [N, 2H].
94
+ eh_input = fused_eh_norm(
95
+ positions,
96
+ inputs_embeds,
97
+ previous_hidden_states,
98
+ self.enorm.weight,
99
+ self.hnorm.weight,
100
+ self.enorm.variance_epsilon,
101
+ )
102
+ hidden_states = self.eh_proj(eh_input)
103
+ # Fuse the residual add and final RMSNorm. Glm5NextMoE already performs
104
+ # its all-reduce, so no collective is needed here. The post-norm result
105
+ # feeds both draft logits and the next recycled hidden state.
106
+ hidden_states, residual, _, _ = self.mtp_block(
107
+ positions=positions, hidden_states=hidden_states, residual=None
108
+ )
109
+ hidden_states, _ = self.shared_head.norm(hidden_states, residual=residual)
110
+ return hidden_states, hidden_states
111
+
112
+
113
+class Glm5NextMultiTokenPredictor(nn.Module):
114
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
115
+ super().__init__()
116
+ config = vllm_config.model_config.hf_config
117
+ self.mtp_start_layer_idx = config.num_hidden_layers
118
+ self.num_mtp_layers = config.num_nextn_predict_layers
119
+ self.layers = torch.nn.ModuleDict(
120
+ {
121
+ str(idx): Glm5NextMultiTokenPredictorLayer(
122
+ vllm_config, f"{prefix}.layers.{idx}"
123
+ )
124
+ for idx in range(
125
+ self.mtp_start_layer_idx,
126
+ self.mtp_start_layer_idx + self.num_mtp_layers,
127
+ )
128
+ }
129
+ )
130
+ self.embed_tokens = VocabParallelEmbedding(
131
+ config.vocab_size,
132
+ config.hidden_size,
133
+ prefix=maybe_prefix(prefix, "embed_tokens"),
134
+ )
135
+ # Plain list for the per-propose lookup: ModuleDict[str(...)] builds a
136
+ # string and hashes it on every draft step.
137
+ self._mtp_layers = list(self.layers.values())
138
+ self._mtp_mla_attns = []
139
+ for layer in self._mtp_layers:
140
+ self_attn = layer.mtp_block.self_attn
141
+ assert isinstance(self_attn, Glm5NextMLAAttention)
142
+ self._mtp_mla_attns.append(self_attn.mla_attn)
143
+ self.logits_processor = LogitsProcessor(config.vocab_size)
144
+
145
+ def set_skip_topk(self, skip: bool):
146
+ # index_share_for_mtp_iteration: step 0 computes top-k, steps 1+ reuse.
147
+ for mla_attn in self._mtp_mla_attns:
148
+ mla_attn.skip_topk = skip
149
+
150
+ def compact_topk_indices(self, slot_ids: torch.Tensor):
151
+ """Gather the top-k index rows at ``slot_ids`` to the front of the buffer."""
152
+ num_slots = slot_ids.numel()
153
+ for mla_attn in self._mtp_mla_attns:
154
+ topk_indices_buffer = mla_attn.topk_indices_buffer
155
+ assert topk_indices_buffer is not None
156
+ topk_indices_buffer[:num_slots] = topk_indices_buffer[slot_ids]
157
+
158
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
159
+ return self.embed_tokens(input_ids)
160
+
161
+ def forward(
162
+ self,
163
+ input_ids: torch.Tensor,
164
+ positions: torch.Tensor,
165
+ previous_hidden_states: torch.Tensor,
166
+ inputs_embeds: torch.Tensor | None = None,
167
+ spec_step_idx: int = 0,
168
+ ) -> torch.Tensor:
169
+ if inputs_embeds is None:
170
+ inputs_embeds = self.embed_tokens(input_ids)
171
+ current_step_idx = spec_step_idx % self.num_mtp_layers
172
+ return self._mtp_layers[current_step_idx](
173
+ input_ids,
174
+ positions,
175
+ previous_hidden_states,
176
+ inputs_embeds,
177
+ current_step_idx,
178
+ )
179
+
180
+ def compute_logits(
181
+ self,
182
+ hidden_states: torch.Tensor,
183
+ spec_step_idx: int = 0,
184
+ ) -> torch.Tensor:
185
+ current_step_idx = spec_step_idx % self.num_mtp_layers
186
+ mtp_layer = self._mtp_layers[current_step_idx]
187
+ # hidden_states is already post-final-norm (produced in the layer
188
+ # forward and recycled as-is); apply the LM head only, without a
189
+ # second RMSNorm.
190
+ return self.logits_processor(mtp_layer.shared_head.head, hidden_states)
191
+
192
+ def get_top_tokens(
193
+ self,
194
+ hidden_states: torch.Tensor,
195
+ spec_step_idx: int = 0,
196
+ ) -> torch.Tensor:
197
+ current_step_idx = spec_step_idx % self.num_mtp_layers
198
+ mtp_layer = self._mtp_layers[current_step_idx]
199
+ # Vocab-parallel argmax for the greedy draft: per-rank head projection
200
+ # + local argmax + a [batch, 2*tp] (value, index) reduce, instead of
201
+ # materializing and all-gathering full [N, vocab] logits per draft
202
+ # step. Tie-breaking matches the full argmax (shards are contiguous
203
+ # and rank-ordered, so the lowest-rank winner is the lowest global
204
+ # index), so greedy draft tokens are unchanged.
205
+ return self.logits_processor.get_top_tokens(
206
+ mtp_layer.shared_head.head, hidden_states
207
+ )
208
+
209
+
210
+class Glm5NextMTP(nn.Module, DeepseekV2MixtureOfExperts):
211
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
212
+ super().__init__()
213
+ self.config = vllm_config.model_config.hf_config
214
+ self.quant_config = vllm_config.quant_config
215
+ self.model = Glm5NextMultiTokenPredictor(
216
+ vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
217
+ )
218
+ self.set_moe_parameters()
219
+
220
+ def set_moe_parameters(self):
221
+ self.num_moe_layers = self.config.num_nextn_predict_layers
222
+ self.num_expert_groups = self.config.n_group
223
+ self.moe_layers = []
224
+ self.moe_mlp_layers = []
225
+ example_moe = None
226
+ for layer in self.model.layers.values():
227
+ mlp = layer.mtp_block.mlp
228
+ if isinstance(mlp, Glm5NextMoE):
229
+ example_moe = mlp
230
+ self.moe_mlp_layers.append(mlp)
231
+ self.moe_layers.append(mlp.experts)
232
+ self.extract_moe_parameters(example_moe)
233
+
234
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
235
+ return self.model.embed_input_ids(input_ids)
236
+
237
+ def forward(
238
+ self,
239
+ input_ids: torch.Tensor | None,
240
+ positions: torch.Tensor,
241
+ hidden_states: torch.Tensor,
242
+ intermediate_tensors: IntermediateTensors | None = None,
243
+ inputs_embeds: torch.Tensor | None = None,
244
+ spec_step_idx: int = 0,
245
+ ) -> torch.Tensor:
246
+ return self.model(
247
+ input_ids, positions, hidden_states, inputs_embeds, spec_step_idx
248
+ )
249
+
250
+ def compute_logits(
251
+ self,
252
+ hidden_states: torch.Tensor,
253
+ spec_step_idx: int = 0,
254
+ ) -> torch.Tensor | None:
255
+ return self.model.compute_logits(hidden_states, spec_step_idx)
256
+
257
+ def get_top_tokens(
258
+ self,
259
+ hidden_states: torch.Tensor,
260
+ spec_step_idx: int = 0,
261
+ ) -> torch.Tensor:
262
+ # Greedy-draft path used when use_local_argmax_reduction is enabled:
263
+ # vocab-parallel argmax, no full-vocab logits.
264
+ return self.model.get_top_tokens(hidden_states, spec_step_idx)
265
+
266
+ def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
267
+ spec_layer_weight_names = [
268
+ "embed_tokens",
269
+ "enorm",
270
+ "hnorm",
271
+ "eh_proj",
272
+ "shared_head",
273
+ ]
274
+ shared_weight_names = ["embed_tokens"]
275
+ spec_layer_weight = False
276
+ shared_weight = False
277
+ for weight_name in spec_layer_weight_names:
278
+ if weight_name in name:
279
+ spec_layer_weight = True
280
+ if weight_name in shared_weight_names:
281
+ shared_weight = True
282
+ break
283
+ if not spec_layer_weight:
284
+ name = name.replace(
285
+ f"model.layers.{spec_layer}.", f"model.layers.{spec_layer}.mtp_block."
286
+ )
287
+ elif shared_weight:
288
+ name = name.replace(f"model.layers.{spec_layer}.", "model.")
289
+ return name
290
+
291
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
292
+ stacked_params_mapping = [
293
+ ("gate_up_proj", "gate_proj", 0),
294
+ ("gate_up_proj", "up_proj", 1),
295
+ ("fused_qkv_a_proj", "q_a_proj", 0),
296
+ ("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1),
297
+ ("wk_weights_proj", "wk", 0),
298
+ ("wk_weights_proj", "weights_proj", 1),
299
+ ]
300
+ expert_params_mapping = fused_moe_make_expert_params_mapping(
301
+ self,
302
+ ckpt_gate_proj_name="gate_proj",
303
+ ckpt_down_proj_name="down_proj",
304
+ ckpt_up_proj_name="up_proj",
305
+ num_experts=self.config.n_routed_experts,
306
+ )
307
+
308
+ params_dict = dict(self.named_parameters())
309
+ loaded_params: set[str] = set()
310
+ _pending_wk_fp8: dict = {}
311
+ # GLM-5.3-Flash NoPE checkpoints omit the RoPE rows from
312
+ # ``kv_a_proj_with_mqa``; the FP8-to-BF16 path pads them for the model.
313
+ kv_a_pad_size = 0
314
+ if self.config.mla_nope and self.config.qk_rope_head_dim > 0:
315
+ kv_a_pad_size = self.config.qk_rope_head_dim
316
+ for name, loaded_weight in weights:
317
+ if "rotary_emb.inv_freq" in name:
318
+ continue
319
+ # Multimodal (Glm5NextForConditionalGeneration) checkpoints prefix
320
+ # the text-tower weights with "model.language_model."; the MTP head
321
+ # is built as a text-only model (model.layers.*), so strip the
322
+ # prefix to match.
323
+ if name.startswith("model.language_model."):
324
+ name = name.replace("model.language_model.", "model.", 1)
325
+ spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
326
+ if spec_layer is None:
327
+ continue
328
+ name = self._rewrite_spec_layer_name(spec_layer, name)
329
+
330
+ if _try_load_fp8_indexer_wk(
331
+ name,
332
+ loaded_weight,
333
+ _pending_wk_fp8,
334
+ params_dict,
335
+ loaded_params,
336
+ ):
337
+ continue
338
+
339
+ # FP8 checkpoint: dequantize the BF16-kept MLA projections
340
+ # (q_a_proj / kv_a_proj_with_mqa / o_proj) to BF16, mirroring the
341
+ # target model. The model holds fused_qkv_a_proj / o_proj in BF16,
342
+ # so the checkpoint's block-FP8 weight + weight_scale_inv for these
343
+ # has no param home and would KeyError without this dequant.
344
+ if _try_load_fp8_attn_proj(
345
+ name,
346
+ loaded_weight,
347
+ _pending_wk_fp8,
348
+ params_dict,
349
+ loaded_params,
350
+ kv_a_pad_size,
351
+ ):
352
+ continue
353
+
354
+ for param_name, weight_name, shard_id in stacked_params_mapping:
355
+ if weight_name not in name:
356
+ continue
357
+ if ("mlp.experts." in name) and name not in params_dict:
358
+ continue
359
+ name_mapped = name.replace(weight_name, param_name)
360
+ if (
361
+ param_name == "fused_qkv_a_proj"
362
+ ) and name_mapped not in params_dict:
363
+ continue
364
+ else:
365
+ name = name_mapped
366
+ if name.endswith(".bias") and name not in params_dict:
367
+ continue
368
+ param = params_dict[name]
369
+ weight_loader = param.weight_loader
370
+ weight_loader(param, loaded_weight, shard_id)
371
+ break
372
+ else:
373
+ is_expert_weight = False
374
+ for mapping in expert_params_mapping:
375
+ param_name, weight_name, expert_id, shard_id = mapping # type: ignore[assignment]
376
+ if weight_name not in name:
377
+ continue
378
+ is_expert_weight = True
379
+ name_mapped = name.replace(weight_name, param_name)
380
+ param = params_dict[name_mapped]
381
+ weight_loader = typing.cast(
382
+ Callable[..., bool], param.weight_loader
383
+ )
384
+ success = weight_loader(
385
+ param,
386
+ loaded_weight,
387
+ name_mapped,
388
+ shard_id=shard_id,
389
+ expert_id=expert_id,
390
+ return_success=True,
391
+ )
392
+ if success:
393
+ name = name_mapped
394
+ break
395
+ else:
396
+ if is_expert_weight:
397
+ continue
398
+ if name.endswith(".bias") and name not in params_dict:
399
+ continue
400
+ name = maybe_remap_kv_scale_name(name, params_dict) # type: ignore[assignment]
401
+ if name is None:
402
+ continue
403
+ if (
404
+ spec_layer != self.model.mtp_start_layer_idx
405
+ and ".layers" not in name
406
+ ):
407
+ continue
408
+ param = params_dict[name]
409
+ weight_loader = getattr(
410
+ param, "weight_loader", default_weight_loader
411
+ )
412
+ weight_loader(param, loaded_weight)
413
+ loaded_params.add(name)
414
+
415
+ loaded_layers: set[int] = set()
416
+ for param_name in loaded_params:
417
+ spec_layer = get_spec_layer_idx_from_weight_name(self.config, param_name)
418
+ if spec_layer is not None:
419
+ loaded_layers.add(spec_layer)
420
+ for layer_idx in range(
421
+ self.model.mtp_start_layer_idx,
422
+ self.model.mtp_start_layer_idx + self.model.num_mtp_layers,
423
+ ):
424
+ if layer_idx not in loaded_layers:
425
+ raise ValueError(
426
+ f"MTP speculative decoding layer {layer_idx} weights "
427
+ f"missing from checkpoint."
428
+ )
429
+ return loaded_params
vllm/models/glm5next/nvidia/multimodal.py
new
+732
@@ -0,0 +1,732 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""GLM-5.3-Flash vision tower and multimodal processor."""
4
+
5
+from collections.abc import Mapping
6
+from functools import cached_property, partial
7
+
8
+import numpy as np
9
+import torch
10
+import torch.nn as nn
11
+from einops import rearrange
12
+
13
+from vllm.distributed import (
14
+ get_tensor_model_parallel_world_size,
15
+ parallel_state,
16
+)
17
+from vllm.distributed import utils as dist_utils
18
+from vllm.model_executor.layers.activation import SiluAndMulWithClamp
19
+from vllm.model_executor.layers.attention import MMEncoderAttention
20
+from vllm.model_executor.layers.conv import Conv2dLayer, Conv3dLayer
21
+from vllm.model_executor.layers.layernorm import RMSNorm
22
+from vllm.model_executor.layers.linear import (
23
+ ColumnParallelLinear,
24
+ MergedColumnParallelLinear,
25
+ QKVParallelLinear,
26
+ RowParallelLinear,
27
+)
28
+from vllm.model_executor.layers.quantization import QuantizationConfig
29
+from vllm.model_executor.layers.rotary_embedding import get_rope
30
+from vllm.model_executor.layers.rotary_embedding.common import ApplyRotaryEmb
31
+from vllm.model_executor.models.glm4_1v import (
32
+ Glm4vMultiModalProcessor,
33
+ Glm4vProcessingInfo,
34
+)
35
+from vllm.model_executor.models.utils import (
36
+ AutoWeightsLoader,
37
+ WeightsMapper,
38
+)
39
+from vllm.model_executor.models.vision import (
40
+ get_vit_attn_backend,
41
+ is_vit_use_data_parallel,
42
+)
43
+from vllm.models.common.ops import fused_q_kv_rmsnorm
44
+from vllm.multimodal.parse import ImageSize, MultiModalDataItems
45
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
46
+
47
+
48
+class Glm5NextVisionPatchEmbed(nn.Module):
49
+ def __init__(
50
+ self,
51
+ patch_size: int = 14,
52
+ temporal_patch_size: int = 1,
53
+ in_channels: int = 3,
54
+ hidden_size: int = 1536,
55
+ ) -> None:
56
+ super().__init__()
57
+ self.patch_size = patch_size
58
+ self.temporal_patch_size = temporal_patch_size
59
+ self.hidden_size = hidden_size
60
+
61
+ kernel_size = (temporal_patch_size, patch_size, patch_size)
62
+ self.proj = Conv3dLayer(
63
+ in_channels,
64
+ hidden_size,
65
+ kernel_size=kernel_size,
66
+ stride=kernel_size,
67
+ bias=True,
68
+ )
69
+
70
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
71
+ L, C = x.shape
72
+ x = x.view(L, -1, self.temporal_patch_size, self.patch_size, self.patch_size)
73
+ x = self.proj(x).view(L, self.hidden_size)
74
+ return x
75
+
76
+
77
+class Glm5NextVisionMLP(nn.Module):
78
+ def __init__(
79
+ self,
80
+ in_features: int,
81
+ hidden_features: int,
82
+ swiglu_limit: float,
83
+ bias: bool = True,
84
+ quant_config: QuantizationConfig | None = None,
85
+ prefix: str = "",
86
+ ):
87
+ super().__init__()
88
+ use_data_parallel = is_vit_use_data_parallel()
89
+ self.gate_up_proj = MergedColumnParallelLinear(
90
+ input_size=in_features,
91
+ output_sizes=[hidden_features] * 2,
92
+ bias=bias,
93
+ quant_config=quant_config,
94
+ prefix=f"{prefix}.gate_up_proj",
95
+ disable_tp=use_data_parallel,
96
+ )
97
+ self.down_proj = RowParallelLinear(
98
+ hidden_features,
99
+ in_features,
100
+ bias=bias,
101
+ quant_config=quant_config,
102
+ prefix=f"{prefix}.down_proj",
103
+ disable_tp=use_data_parallel,
104
+ )
105
+ # GLM-5.3-Flash clamps the vision SwiGLU gate/up unlike GLM-OCR/GLM-4V.
106
+ self.act_fn = SiluAndMulWithClamp(swiglu_limit=swiglu_limit)
107
+
108
+ def forward(self, x: torch.Tensor):
109
+ x, _ = self.gate_up_proj(x)
110
+ x = self.act_fn(x)
111
+ x, _ = self.down_proj(x)
112
+ return x
113
+
114
+
115
+class Glm5NextVisionAttention(nn.Module):
116
+ def __init__(
117
+ self,
118
+ embed_dim: int,
119
+ num_heads: int,
120
+ projection_size: int,
121
+ quant_config: QuantizationConfig | None = None,
122
+ prefix: str = "",
123
+ ) -> None:
124
+ super().__init__()
125
+ use_data_parallel = is_vit_use_data_parallel()
126
+ self.tp_size = (
127
+ 1 if use_data_parallel else get_tensor_model_parallel_world_size()
128
+ )
129
+ self.tp_rank = (
130
+ 0 if use_data_parallel else parallel_state.get_tensor_model_parallel_rank()
131
+ )
132
+ self.hidden_size_per_attention_head = dist_utils.divide(
133
+ projection_size, num_heads
134
+ )
135
+ self.num_attention_heads_per_partition = dist_utils.divide(
136
+ num_heads, self.tp_size
137
+ )
138
+
139
+ self.head_dim = embed_dim // num_heads
140
+
141
+ # q/k norm eps hard-coded 1e-5 — distinct from block/post norm eps.
142
+ self.q_norm = RMSNorm(self.head_dim, eps=1e-5)
143
+ self.k_norm = RMSNorm(self.head_dim, eps=1e-5)
144
+
145
+ self.qkv = QKVParallelLinear(
146
+ hidden_size=embed_dim,
147
+ head_size=self.hidden_size_per_attention_head,
148
+ total_num_heads=num_heads,
149
+ total_num_kv_heads=num_heads,
150
+ bias=True,
151
+ quant_config=quant_config,
152
+ prefix=f"{prefix}.qkv_proj" if quant_config else f"{prefix}.qkv",
153
+ disable_tp=use_data_parallel,
154
+ )
155
+ self.proj = RowParallelLinear(
156
+ input_size=projection_size,
157
+ output_size=embed_dim,
158
+ quant_config=quant_config,
159
+ prefix=f"{prefix}.proj",
160
+ bias=True,
161
+ disable_tp=use_data_parallel,
162
+ )
163
+
164
+ self.attn = MMEncoderAttention(
165
+ num_heads=self.num_attention_heads_per_partition,
166
+ head_size=self.hidden_size_per_attention_head,
167
+ scale=self.hidden_size_per_attention_head**-0.5,
168
+ prefix=f"{prefix}.attn",
169
+ )
170
+ self.apply_rotary_emb = ApplyRotaryEmb(enforce_enable=True)
171
+
172
+ def split_qkv(self, qkv: torch.Tensor) -> tuple[torch.Tensor, ...]:
173
+ seq_len, bs, _ = qkv.shape
174
+ q, k, v = qkv.chunk(3, dim=2)
175
+ new_shape = (
176
+ seq_len,
177
+ bs,
178
+ self.num_attention_heads_per_partition,
179
+ self.hidden_size_per_attention_head,
180
+ )
181
+ q, k, v = (x.view(*new_shape) for x in (q, k, v))
182
+ return q, k, v
183
+
184
+ def forward(
185
+ self,
186
+ x: torch.Tensor,
187
+ cu_seqlens: torch.Tensor,
188
+ rotary_pos_emb_cos: torch.Tensor,
189
+ rotary_pos_emb_sin: torch.Tensor,
190
+ max_seqlen: torch.Tensor | None = None,
191
+ ) -> torch.Tensor:
192
+ x, _ = self.qkv(x)
193
+ q, k, v = self.split_qkv(x)
194
+
195
+ # P1: fused q/k RMSNorm (two distinct weights, one launch; fp32, bit-identical).
196
+ q_shape, k_shape = q.shape, k.shape
197
+ q_flat = q.reshape(-1, self.head_dim)
198
+ k_flat = k.reshape(-1, self.head_dim)
199
+ q, k = fused_q_kv_rmsnorm(
200
+ q_flat,
201
+ k_flat,
202
+ self.q_norm.weight,
203
+ self.k_norm.weight,
204
+ self.q_norm.variance_epsilon,
205
+ )
206
+ q = q.view(q_shape)
207
+ k = k.view(k_shape)
208
+
209
+ q, k, v = (rearrange(t, "s b ... -> b s ...").contiguous() for t in (q, k, v))
210
+ if rotary_pos_emb_cos is not None and rotary_pos_emb_sin is not None:
211
+ qk_concat = torch.cat([q, k], dim=0)
212
+ qk_rotated = self.apply_rotary_emb(
213
+ qk_concat,
214
+ rotary_pos_emb_cos,
215
+ rotary_pos_emb_sin,
216
+ )
217
+ q, k = torch.chunk(qk_rotated, 2, dim=0)
218
+
219
+ context_layer = self.attn(
220
+ query=q,
221
+ key=k,
222
+ value=v,
223
+ cu_seqlens=cu_seqlens,
224
+ max_seqlen=max_seqlen,
225
+ )
226
+ context_layer = rearrange(context_layer, "b s h d -> s b (h d)").contiguous()
227
+
228
+ output, _ = self.proj(context_layer)
229
+ return output
230
+
231
+
232
+class Glm5NextVisionBlock(nn.Module):
233
+ def __init__(
234
+ self,
235
+ dim: int,
236
+ num_heads: int,
237
+ mlp_hidden_dim: int,
238
+ swiglu_limit: float,
239
+ norm_layer: partial[nn.Module] | None = None,
240
+ quant_config: QuantizationConfig | None = None,
241
+ prefix: str = "",
242
+ ) -> None:
243
+ super().__init__()
244
+ if norm_layer is None:
245
+ norm_layer = partial(nn.LayerNorm, eps=1e-6)
246
+ self.norm1 = norm_layer(dim)
247
+ self.norm2 = norm_layer(dim)
248
+ self.attn = Glm5NextVisionAttention(
249
+ embed_dim=dim,
250
+ num_heads=num_heads,
251
+ projection_size=dim,
252
+ quant_config=quant_config,
253
+ prefix=f"{prefix}.attn",
254
+ )
255
+ self.mlp = Glm5NextVisionMLP(
256
+ dim,
257
+ mlp_hidden_dim,
258
+ swiglu_limit=swiglu_limit,
259
+ bias=True,
260
+ quant_config=quant_config,
261
+ prefix=f"{prefix}.mlp",
262
+ )
263
+
264
+ def forward(
265
+ self,
266
+ x: torch.Tensor,
267
+ cu_seqlens: torch.Tensor,
268
+ rotary_pos_emb_cos: torch.Tensor,
269
+ rotary_pos_emb_sin: torch.Tensor,
270
+ max_seqlen: int | None = None,
271
+ ) -> torch.Tensor:
272
+ x_attn = self.attn(
273
+ self.norm1(x),
274
+ cu_seqlens=cu_seqlens,
275
+ rotary_pos_emb_cos=rotary_pos_emb_cos,
276
+ rotary_pos_emb_sin=rotary_pos_emb_sin,
277
+ max_seqlen=max_seqlen,
278
+ )
279
+ x_fused_norm, residual = self.norm2(x, residual=x_attn)
280
+ x = residual + self.mlp(x_fused_norm)
281
+ return x
282
+
283
+
284
+class Glm5NextPatchMerger(nn.Module):
285
+ def __init__(
286
+ self,
287
+ d_model: int,
288
+ context_dim: int,
289
+ swiglu_limit: float,
290
+ quant_config: QuantizationConfig | None = None,
291
+ bias: bool = False,
292
+ prefix: str = "",
293
+ ) -> None:
294
+ super().__init__()
295
+ use_data_parallel = is_vit_use_data_parallel()
296
+ self.hidden_size = d_model
297
+ self.proj = ColumnParallelLinear(
298
+ self.hidden_size,
299
+ self.hidden_size,
300
+ bias=bias,
301
+ gather_output=True,
302
+ quant_config=quant_config,
303
+ prefix=f"{prefix}.proj",
304
+ disable_tp=use_data_parallel,
305
+ )
306
+ self.post_projection_norm = nn.LayerNorm(self.hidden_size)
307
+ self.gate_up_proj = MergedColumnParallelLinear(
308
+ input_size=self.hidden_size,
309
+ output_sizes=[context_dim] * 2,
310
+ bias=bias,
311
+ quant_config=quant_config,
312
+ prefix=f"{prefix}.gate_up_proj",
313
+ disable_tp=use_data_parallel,
314
+ )
315
+ self.down_proj = RowParallelLinear(
316
+ context_dim,
317
+ self.hidden_size,
318
+ bias=bias,
319
+ quant_config=quant_config,
320
+ prefix=f"{prefix}.down_proj",
321
+ disable_tp=use_data_parallel,
322
+ )
323
+ # GLM-5.3-Flash also clamps the merger SwiGLU.
324
+ self.act_fn = SiluAndMulWithClamp(swiglu_limit=swiglu_limit)
325
+ self.extra_activation_func = nn.GELU()
326
+
327
+ def forward(self, x: torch.Tensor):
328
+ x, _ = self.proj(x)
329
+ x = self.extra_activation_func(self.post_projection_norm(x))
330
+ gate_up, _ = self.gate_up_proj(x)
331
+ x = self.act_fn(gate_up)
332
+ x, _ = self.down_proj(x)
333
+ return x
334
+
335
+
336
+class Glm5NextVisionTransformer(nn.Module):
337
+ # Stacked-weight remap for the GLM-OCR/GLM-4V vision checkpoint layout.
338
+ hf_to_vllm_mapper = WeightsMapper(
339
+ orig_to_new_stacked={
340
+ ".attn.q.": (".attn.qkv.", "q"),
341
+ ".attn.k.": (".attn.qkv.", "k"),
342
+ ".attn.v.": (".attn.qkv.", "v"),
343
+ ".gate_proj": (".gate_up_proj", 0),
344
+ ".up_proj": (".gate_up_proj", 1),
345
+ }
346
+ )
347
+
348
+ def __init__(
349
+ self,
350
+ text_config, # noqa: ANN001
351
+ vision_config,
352
+ norm_eps: float = 1e-6,
353
+ quant_config: QuantizationConfig | None = None,
354
+ prefix: str = "",
355
+ ) -> None:
356
+ super().__init__()
357
+ use_data_parallel = is_vit_use_data_parallel()
358
+ self.tp_size = (
359
+ 1 if use_data_parallel else get_tensor_model_parallel_world_size()
360
+ )
361
+
362
+ patch_size = vision_config.patch_size
363
+ temporal_patch_size = vision_config.temporal_patch_size
364
+ in_channels = vision_config.in_channels
365
+ depth = vision_config.depth
366
+ self.hidden_size = vision_config.hidden_size
367
+ self.num_heads = vision_config.num_heads
368
+
369
+ self.patch_size = vision_config.patch_size
370
+ self.spatial_merge_size = vision_config.spatial_merge_size
371
+ self.out_hidden_size = vision_config.out_hidden_size
372
+
373
+ swiglu_limit = vision_config.swiglu_limit
374
+ if swiglu_limit is None:
375
+ swiglu_limit = text_config.swiglu_limit
376
+ assert swiglu_limit is not None, (
377
+ "GLM-5.3-Flash vision requires swiglu_limit (vision_config or text_config)"
378
+ )
379
+
380
+ # Single construction pass — no abs-pos embeddings / post-conv norm (OCR delta).
381
+ self.patch_embed = Glm5NextVisionPatchEmbed(
382
+ patch_size=patch_size,
383
+ temporal_patch_size=temporal_patch_size,
384
+ in_channels=in_channels,
385
+ hidden_size=self.hidden_size,
386
+ )
387
+
388
+ norm_layer = partial(RMSNorm, eps=norm_eps)
389
+ head_dim = self.hidden_size // self.num_heads
390
+ self.rotary_pos_emb = get_rope(
391
+ head_size=head_dim,
392
+ max_position=8192,
393
+ is_neox_style=True,
394
+ rope_parameters={"partial_rotary_factor": 0.5},
395
+ )
396
+ self.blocks = nn.ModuleList(
397
+ [
398
+ Glm5NextVisionBlock(
399
+ dim=self.hidden_size,
400
+ num_heads=self.num_heads,
401
+ mlp_hidden_dim=vision_config.intermediate_size,
402
+ swiglu_limit=swiglu_limit,
403
+ norm_layer=norm_layer,
404
+ quant_config=quant_config,
405
+ prefix=f"{prefix}.blocks.{layer_idx}",
406
+ )
407
+ for layer_idx in range(depth)
408
+ ]
409
+ )
410
+ # GLM-5.3-Flash merger bottleneck width.
411
+ self.merger = Glm5NextPatchMerger(
412
+ d_model=vision_config.out_hidden_size,
413
+ context_dim=vision_config.projection_intermediate_size,
414
+ swiglu_limit=swiglu_limit,
415
+ quant_config=quant_config,
416
+ bias=False,
417
+ prefix=f"{prefix}.merger",
418
+ )
419
+
420
+ self.downsample = Conv2dLayer(
421
+ in_channels=vision_config.hidden_size,
422
+ out_channels=vision_config.out_hidden_size,
423
+ kernel_size=vision_config.spatial_merge_size,
424
+ stride=vision_config.spatial_merge_size,
425
+ )
426
+ self.post_layernorm = RMSNorm(
427
+ vision_config.hidden_size, eps=vision_config.rms_norm_eps
428
+ )
429
+
430
+ self.attn_backend = get_vit_attn_backend(
431
+ head_size=head_dim,
432
+ dtype=torch.get_default_dtype(),
433
+ )
434
+
435
+ @property
436
+ def dtype(self) -> torch.dtype:
437
+ return self.patch_embed.proj.weight.dtype
438
+
439
+ @property
440
+ def device(self) -> torch.device:
441
+ return self.patch_embed.proj.weight.device
442
+
443
+ def rot_pos_emb(
444
+ self, grid_thw: list[list[int]]
445
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
446
+ pos_ids = []
447
+ for t, h, w in grid_thw:
448
+ hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w)
449
+ wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)
450
+ hpos_ids = (
451
+ hpos_ids.reshape(
452
+ h // self.spatial_merge_size,
453
+ self.spatial_merge_size,
454
+ w // self.spatial_merge_size,
455
+ self.spatial_merge_size,
456
+ )
457
+ .permute(0, 2, 1, 3)
458
+ .flatten()
459
+ )
460
+ wpos_ids = (
461
+ wpos_ids.reshape(
462
+ h // self.spatial_merge_size,
463
+ self.spatial_merge_size,
464
+ w // self.spatial_merge_size,
465
+ self.spatial_merge_size,
466
+ )
467
+ .permute(0, 2, 1, 3)
468
+ .flatten()
469
+ )
470
+ pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1))
471
+ pos_ids = torch.cat(pos_ids, dim=0)
472
+ max_grid_size = max(max(h, w) for _, h, w in grid_thw)
473
+
474
+ cos, sin = self.rotary_pos_emb.get_cos_sin(max_grid_size)
475
+
476
+ pos_ids = pos_ids.to(cos.device, non_blocking=True)
477
+ cos_combined = cos[pos_ids].flatten(1)
478
+ sin_combined = sin[pos_ids].flatten(1)
479
+ return cos_combined, sin_combined, pos_ids
480
+
481
+ def compute_attn_mask_seqlen(
482
+ self,
483
+ cu_seqlens: torch.Tensor,
484
+ ) -> torch.Tensor | None:
485
+ max_seqlen = None
486
+ if self.attn_backend in {
487
+ AttentionBackendEnum.FLASH_ATTN,
488
+ AttentionBackendEnum.ROCM_AITER_FA,
489
+ AttentionBackendEnum.TRITON_ATTN,
490
+ }:
491
+ max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
492
+ return max_seqlen
493
+
494
+ def prepare_encoder_metadata(
495
+ self,
496
+ grid_thw_list: list[list[int]],
497
+ *,
498
+ max_batch_size: int | None = None,
499
+ max_frames_per_batch: int | None = None,
500
+ max_seqlen_override: int | None = None,
501
+ device: torch.device | None = None,
502
+ ) -> dict[str, torch.Tensor | None]:
503
+ """Compute encoder metadata for eager and CUDA graph execution."""
504
+ if device is None:
505
+ device = self.device
506
+
507
+ metadata: dict[str, torch.Tensor | None] = {}
508
+
509
+ rotary_cos, rotary_sin, _ = self.rot_pos_emb(grid_thw_list)
510
+ metadata["rotary_pos_emb_cos"] = rotary_cos
511
+ metadata["rotary_pos_emb_sin"] = rotary_sin
512
+
513
+ grid_thw_np = np.array(grid_thw_list, dtype=np.int32)
514
+ patches_per_frame = grid_thw_np[:, 1] * grid_thw_np[:, 2]
515
+ cu_seqlens = np.repeat(patches_per_frame, grid_thw_np[:, 0]).cumsum(
516
+ dtype=np.int32
517
+ )
518
+ cu_seqlens = np.concatenate([np.zeros(1, dtype=np.int32), cu_seqlens])
519
+
520
+ pad_to = (
521
+ max_frames_per_batch if max_frames_per_batch is not None else max_batch_size
522
+ )
523
+ if pad_to is not None:
524
+ num_seqs = len(cu_seqlens) - 1
525
+ if num_seqs < pad_to:
526
+ cu_seqlens = np.concatenate(
527
+ [
528
+ cu_seqlens,
529
+ np.full(
530
+ pad_to - num_seqs,
531
+ cu_seqlens[-1],
532
+ dtype=np.int32,
533
+ ),
534
+ ]
535
+ )
536
+
537
+ metadata["sequence_lengths"] = MMEncoderAttention.maybe_compute_seq_lens(
538
+ self.attn_backend, cu_seqlens, device
539
+ )
540
+
541
+ if max_seqlen_override is not None:
542
+ max_seqlen_val = max_seqlen_override
543
+ else:
544
+ max_seqlen_val = MMEncoderAttention.compute_max_seqlen(
545
+ self.attn_backend, cu_seqlens
546
+ )
547
+ metadata["max_seqlen"] = torch.tensor(max_seqlen_val, dtype=torch.int32)
548
+
549
+ metadata["cu_seqlens"] = MMEncoderAttention.maybe_recompute_cu_seqlens(
550
+ self.attn_backend,
551
+ cu_seqlens,
552
+ self.hidden_size,
553
+ self.tp_size,
554
+ device,
555
+ )
556
+
557
+ return metadata
558
+
559
+ def forward(
560
+ self,
561
+ x: torch.Tensor,
562
+ grid_thw: torch.Tensor | list[list[int]],
563
+ *,
564
+ encoder_metadata: dict[str, torch.Tensor] | None = None,
565
+ ) -> torch.Tensor:
566
+ # patchify
567
+ x = x.to(device=self.device, dtype=self.dtype)
568
+ x = self.patch_embed(x)
569
+
570
+ if encoder_metadata is not None:
571
+ # Encoder CUDA-graph path (PR #49852): rotary/cu_seqlens/max_seqlen are
572
+ # precomputed by prepare_encoder_metadata (which uses rot_pos_emb exactly
573
+ # as the eager rebuild does), so reuse them and skip the per-call CPU
574
+ # rebuild (the low-GPU-util culprit on multimodal workloads).
575
+ rotary_pos_emb_cos = encoder_metadata["rotary_pos_emb_cos"]
576
+ rotary_pos_emb_sin = encoder_metadata["rotary_pos_emb_sin"]
577
+ cu_seqlens = encoder_metadata["cu_seqlens"]
578
+ max_seqlen = encoder_metadata["max_seqlen"]
579
+ else:
580
+ if isinstance(grid_thw, list):
581
+ grid_thw = torch.tensor(grid_thw, dtype=torch.int32)
582
+ rotary_pos_emb_cos, rotary_pos_emb_sin, _ = self.rot_pos_emb(grid_thw)
583
+ cu_seqlens = torch.repeat_interleave(
584
+ grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]
585
+ ).cumsum(dim=0, dtype=torch.int32)
586
+ cu_seqlens = torch.cat([cu_seqlens.new_zeros(1), cu_seqlens])
587
+ cu_seqlens = cu_seqlens.to(self.device, non_blocking=True)
588
+ max_seqlen = self.compute_attn_mask_seqlen(cu_seqlens)
589
+
590
+ # transformers
591
+ x = x.unsqueeze(1)
592
+ for blk in self.blocks:
593
+ x = blk(
594
+ x,
595
+ cu_seqlens=cu_seqlens,
596
+ rotary_pos_emb_cos=rotary_pos_emb_cos,
597
+ rotary_pos_emb_sin=rotary_pos_emb_sin,
598
+ max_seqlen=max_seqlen,
599
+ )
600
+
601
+ # adapter
602
+ x = self.post_layernorm(x)
603
+ x = x.view(-1, self.spatial_merge_size, self.spatial_merge_size, x.shape[-1])
604
+ x = x.permute(0, 3, 1, 2)
605
+ x = self.downsample(x).view(-1, self.out_hidden_size)
606
+ x = self.merger(x)
607
+ return x
608
+
609
+ def load_weights(self, weights) -> set[str]:
610
+ loader = AutoWeightsLoader(self)
611
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
612
+
613
+
614
+class Glm5NextProcessingInfo(Glm4vProcessingInfo):
615
+ """Wires up the vLLM-native processor for the multimodal checkpoint.
616
+
617
+ The checkpoint's ``processor_config.json`` declares a custom ``processor_class``
618
+ and stores its image/video processor configs inline (no standalone
619
+ ``preprocessor_config.json``), so ``AutoProcessor`` cannot resolve the
620
+ config. We bypass it and build our own ``Glm5NextProcessor``
621
+ (``vllm/transformers_utils/processors/glm5next.py``), a port of the
622
+ training-side pipeline that no longer imports transformers' GLM processor
623
+ classes. The port applies ``patch_expand_factor`` (checkpoint ships 1)
624
+ inside ``smart_resize``'s spatial factor.
625
+ """
626
+
627
+ @cached_property
628
+ def _glm5_hf_processor(self):
629
+ from vllm.transformers_utils.processors.glm5next import Glm5NextProcessor
630
+
631
+ return Glm5NextProcessor.from_pretrained(self.ctx.model_config.model)
632
+
633
+ def get_hf_processor(self, **kwargs: object):
634
+ return self._glm5_hf_processor
635
+
636
+ def _processor_pixel_budget(self, proc) -> tuple[int, int]:
637
+ from vllm.transformers_utils.processors.glm5next import _pixel_budget
638
+
639
+ return _pixel_budget(
640
+ proc.min_image_tokens,
641
+ proc.max_image_tokens,
642
+ proc.patch_size,
643
+ proc.merge_size,
644
+ proc.temporal_patch_size,
645
+ )
646
+
647
+ def _get_image_max_pixels(self) -> int:
648
+ mm_kwargs = self.ctx.get_merged_mm_kwargs({})
649
+ if (override := mm_kwargs.get("max_pixels")) is not None:
650
+ return int(override)
651
+ return self._processor_pixel_budget(self.get_hf_processor().image_processor)[1]
652
+
653
+ def _get_video_max_pixels(self) -> int:
654
+ mm_kwargs = self.ctx.get_merged_mm_kwargs({})
655
+ if (override := mm_kwargs.get("max_pixels")) is not None:
656
+ return int(override)
657
+ return self._processor_pixel_budget(self.get_hf_processor().video_processor)[1]
658
+
659
+ def _get_vision_info(
660
+ self,
661
+ *,
662
+ image_width: int,
663
+ image_height: int,
664
+ num_frames: int = 16,
665
+ do_resize: bool = True,
666
+ max_image_pixels: int = 28 * 28 * 2 * 30000,
667
+ ) -> tuple[ImageSize, int]:
668
+ """GLM-5.3-Flash canvas geometry for token budgeting and dummy inputs.
669
+
670
+ The inherited Glm4v path resolves the pixel budget from
671
+ ``size.longest_edge`` and resizes with GLM-4V's ``smart_resize``. This
672
+ checkpoint's ``processor_config.json`` ships the token-budget style
673
+ (``min_image_tokens`` / ``max_image_tokens``) with no ``size`` key, and
674
+ the alignment factor carries ``patch_expand_factor`` — resolve both
675
+ from the vLLM-native processor so profiling matches runtime geometry.
676
+ """
677
+ from vllm.transformers_utils.processors.glm5next import smart_resize
678
+
679
+ vision_config = self.get_hf_config().vision_config
680
+ patch_size = vision_config.patch_size
681
+ merge_size = vision_config.spatial_merge_size
682
+ temporal_patch_size = vision_config.temporal_patch_size
683
+
684
+ image_processor = self.get_hf_processor().image_processor
685
+ factor = patch_size * merge_size * image_processor.patch_expand_factor
686
+ # Keep the profiling search viable when the caller's budget is below
687
+ # one aligned canvas of the requested duration.
688
+ max_image_pixels = max(max_image_pixels, temporal_patch_size * factor * factor)
689
+
690
+ if do_resize:
691
+ t = num_frames if num_frames > temporal_patch_size else temporal_patch_size
692
+ resized_height, resized_width = smart_resize(
693
+ t=t,
694
+ h=image_height,
695
+ w=image_width,
696
+ t_factor=temporal_patch_size,
697
+ h_factor=factor,
698
+ w_factor=factor,
699
+ min_pixels=1,
700
+ max_pixels=max_image_pixels,
701
+ )
702
+ preprocessed_size = ImageSize(width=resized_width, height=resized_height)
703
+ else:
704
+ preprocessed_size = ImageSize(width=image_width, height=image_height)
705
+
706
+ padded_num_frames = num_frames + (-num_frames % temporal_patch_size)
707
+ grid_t = max(padded_num_frames // temporal_patch_size, 1)
708
+ grid_h = preprocessed_size.height // patch_size
709
+ grid_w = preprocessed_size.width // patch_size
710
+
711
+ num_patches = grid_t * grid_h * grid_w
712
+ num_vision_tokens = num_patches // (merge_size**2)
713
+
714
+ return preprocessed_size, num_vision_tokens
715
+
716
+
717
+class Glm5NextMultiModalProcessor(Glm4vMultiModalProcessor):
718
+ """The vLLM-native ``Glm5NextProcessor`` extracts image/video features
719
+ only and passes the prompt text through unchanged, so prompt expansion
720
+ (image token repeat, video frame/timestamp structure) is owned by vLLM's
721
+ prompt-update machinery — the inherited ``_get_prompt_updates`` builds
722
+ the replacement content and the placeholder scan validates against
723
+ exactly that."""
724
+
725
+ def _hf_processor_applies_updates(
726
+ self,
727
+ prompt_text: str,
728
+ mm_items: MultiModalDataItems,
729
+ hf_processor_mm_kwargs: Mapping[str, object],
730
+ tokenization_kwargs: Mapping[str, object],
731
+ ) -> bool:
732
+ return False
vllm/models/glm5next/nvidia/ops/fused_eh_norm.py
new
+77
@@ -0,0 +1,77 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import torch
4
+
5
+from vllm.triton_utils import tl, triton
6
+
7
+
8
+@triton.jit
9
+def _rms_norm(x, w, eps, HIDDEN_SIZE: tl.constexpr):
10
+ x = x.to(tl.float32)
11
+ mean_sq = tl.sum(x * x, axis=0) / HIDDEN_SIZE
12
+ rrms = tl.rsqrt(mean_sq + eps)
13
+ w = w.to(tl.float32)
14
+ return (x * rrms) * w
15
+
16
+
17
+@triton.jit
18
+def _fused_eh_norm_kernel(
19
+ pos_ptr,
20
+ embeds_ptr,
21
+ embeds_stride,
22
+ prev_ptr,
23
+ prev_stride,
24
+ enorm_w_ptr,
25
+ hnorm_w_ptr,
26
+ eps,
27
+ out_ptr,
28
+ out_stride,
29
+ H: tl.constexpr,
30
+ BLOCK: tl.constexpr,
31
+):
32
+ """MTP input fusion: zero embeds at position 0, RMSNorm(embeds) with enorm
33
+ and RMSNorm(prev_hidden) with hnorm, written side-by-side into ``out``
34
+ ([N, 2H]) ready for the eh_proj GEMM. Replaces where + 2x RMSNorm + cat."""
35
+ tok = tl.program_id(0)
36
+ off = tl.arange(0, BLOCK)
37
+ mask = off < H
38
+
39
+ pos = tl.load(pos_ptr + tok)
40
+ e = tl.load(embeds_ptr + tok * embeds_stride + off, mask=mask, other=0.0)
41
+ e = tl.where(pos == 0, 0.0, e.to(tl.float32))
42
+ ew = tl.load(enorm_w_ptr + off, mask=mask)
43
+ e_normed = _rms_norm(e, ew, eps, H)
44
+ tl.store(out_ptr + tok * out_stride + off, e_normed, mask=mask)
45
+
46
+ p = tl.load(prev_ptr + tok * prev_stride + off, mask=mask, other=0.0)
47
+ hw = tl.load(hnorm_w_ptr + off, mask=mask)
48
+ p_normed = _rms_norm(p, hw, eps, H)
49
+ tl.store(out_ptr + tok * out_stride + H + off, p_normed, mask=mask)
50
+
51
+
52
+def fused_eh_norm(
53
+ positions: torch.Tensor,
54
+ inputs_embeds: torch.Tensor,
55
+ previous_hidden: torch.Tensor,
56
+ enorm_w: torch.Tensor,
57
+ hnorm_w: torch.Tensor,
58
+ eps: float,
59
+) -> torch.Tensor:
60
+ """Returns cat([enorm(masked embeds), hnorm(prev_hidden)]) -> [N, 2H]."""
61
+ n, h = inputs_embeds.shape
62
+ out = torch.empty(n, 2 * h, dtype=inputs_embeds.dtype, device=inputs_embeds.device)
63
+ _fused_eh_norm_kernel[(n,)](
64
+ positions,
65
+ inputs_embeds,
66
+ inputs_embeds.stride(0),
67
+ previous_hidden,
68
+ previous_hidden.stride(0),
69
+ enorm_w,
70
+ hnorm_w,
71
+ eps,
72
+ out,
73
+ out.stride(0),
74
+ h,
75
+ triton.next_power_of_2(h),
76
+ )
77
+ return out
vllm/models/glm5next/nvidia/ops/kpool_compress.py
new
+891
@@ -0,0 +1,891 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""kpool (key-pooling) Triton kernels for the sparse-attention indexer.
4
+
5
+The cache stores POOLS (1 entry per ``pool_size`` consecutive tokens) rather
6
+than individual tokens. ``compress_ratio == pool_size`` on the kv_cache_spec
7
+makes the metadata builder emit pool-granular slot_mapping / seq_lens /
8
+cu_seq_lens / page_table for free; this file supplies the compress-write
9
+kernel (replacing ``indexer_k_quant_and_cache``) and the pool-level topk
10
+helpers (select pools -> expand to tokens -> append tail).
11
+"""
12
+
13
+from __future__ import annotations
14
+
15
+import torch
16
+
17
+from vllm.triton_utils import tl, triton
18
+
19
+# The GLM-5.3-Flash indexer head dimension is fixed at 128.
20
+INDEX_HEAD_DIM = 128
21
+
22
+
23
+# Hadamard-128 rotation
24
+
25
+
26
+@triton.jit
27
+def _hadamard128_stage(x, GROUPS: tl.constexpr, STRIDE: tl.constexpr):
28
+ x3 = tl.reshape(x, (GROUPS, 2, STRIDE))
29
+ x3 = tl.trans(x3, 0, 2, 1)
30
+ a, b = tl.split(x3)
31
+ x3 = tl.join(a + b, a - b)
32
+ x3 = tl.trans(x3, 0, 2, 1)
33
+ return tl.reshape(x3, (128,))
34
+
35
+
36
+@triton.jit
37
+def _hadamard128(x):
38
+ x = _hadamard128_stage(x, 64, 1)
39
+ x = _hadamard128_stage(x, 32, 2)
40
+ x = _hadamard128_stage(x, 16, 4)
41
+ x = _hadamard128_stage(x, 8, 8)
42
+ x = _hadamard128_stage(x, 4, 16)
43
+ x = _hadamard128_stage(x, 2, 32)
44
+ x = _hadamard128_stage(x, 1, 64)
45
+ return x * 0.08838834764831845 # 1/sqrt(128)
46
+
47
+
48
+@triton.jit
49
+def _fwht_stage(x, N: tl.constexpr, GROUPS: tl.constexpr, STRIDE: tl.constexpr):
50
+ # One FWHT butterfly stage on a flat tensor of N = GROUPS*2*STRIDE elems;
51
+ # same construction as _hadamard128_stage but with a parametric N so it can
52
+ # process BLOCK_R rows at once.
53
+ x3 = tl.reshape(x, (GROUPS, 2, STRIDE))
54
+ x3 = tl.trans(x3, 0, 2, 1)
55
+ a, b = tl.split(x3)
56
+ x3 = tl.join(a + b, a - b)
57
+ x3 = tl.trans(x3, 0, 2, 1)
58
+ return tl.reshape(x3, (N,))
59
+
60
+
61
+@triton.jit
62
+def _fwht_quant_kernel(
63
+ q_ptr,
64
+ qout_ptr,
65
+ sout_ptr,
66
+ n_rows,
67
+ BLOCK_R: tl.constexpr,
68
+):
69
+ """Fused Hadamard-128 rotation + per-row absmax FP8 (ue8m0) quant.
70
+
71
+ Each row uses fp32 butterflies and scaling, rounds to bf16, then applies
72
+ absmax quantization with a power-of-two scale.
73
+ """
74
+ pid = tl.program_id(0)
75
+ rows = pid * BLOCK_R + tl.arange(0, BLOCK_R)
76
+ rmask = rows < n_rows
77
+ offs = tl.arange(0, 128)
78
+ x = tl.load(
79
+ q_ptr + rows[:, None] * 128 + offs[None, :], mask=rmask[:, None], other=0.0
80
+ ).to(tl.float32)
81
+
82
+ # Flatten so each row's 128 lanes stay contiguous: every stage's
83
+ # (GROUPS, 2, STRIDE) tiling has 2*STRIDE dividing 128, so pairs never
84
+ # straddle a row boundary. GROUPS of each stage scales by BLOCK_R.
85
+ N: tl.constexpr = BLOCK_R * 128
86
+ x = tl.reshape(x, (N,))
87
+ x = _fwht_stage(x, N, BLOCK_R * 64, 1)
88
+ x = _fwht_stage(x, N, BLOCK_R * 32, 2)
89
+ x = _fwht_stage(x, N, BLOCK_R * 16, 4)
90
+ x = _fwht_stage(x, N, BLOCK_R * 8, 8)
91
+ x = _fwht_stage(x, N, BLOCK_R * 4, 16)
92
+ x = _fwht_stage(x, N, BLOCK_R * 2, 32)
93
+ x = _fwht_stage(x, N, BLOCK_R, 64)
94
+ x = x * 0.08838834764831845 # 1/sqrt(128), exact in fp32
95
+
96
+ # Match the unfused path's bf16 materialization before quantizing.
97
+ x = x.to(tl.bfloat16).to(tl.float32)
98
+ x = tl.reshape(x, (BLOCK_R, 128))
99
+
100
+ absmax = tl.maximum(tl.max(tl.abs(x), axis=1), 1e-4)
101
+ scale = tl.exp2(tl.ceil(tl.log2(absmax * (1.0 / 448.0))))
102
+ y = tl.minimum(tl.maximum(x / scale[:, None], -448.0), 448.0)
103
+
104
+ tl.store(qout_ptr + rows[:, None] * 128 + offs[None, :], y, mask=rmask[:, None])
105
+ tl.store(sout_ptr + rows, scale, mask=rmask)
106
+
107
+
108
+def fwht128_quant_fp8(q: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
109
+ """Rotate each 128-wide row by the Hadamard-128 transform, then FP8-quant.
110
+
111
+ The fused kernel avoids materializing the rotated tensor and uses an exact
112
+ fp32 ``1 / sqrt(128)`` scale before block-128 ue8m0 quantization.
113
+
114
+ Args:
115
+ q: ``[rows, 128]`` bf16 — one head vector per row.
116
+
117
+ Returns:
118
+ (q_fp8 ``[rows, 128]`` float8_e4m3fn, scale ``[rows, 1]`` float32).
119
+ """
120
+ assert q.ndim == 2 and q.shape[1] == 128, q.shape
121
+ assert q.dtype == torch.bfloat16
122
+ assert q.is_contiguous()
123
+ n_rows = q.shape[0]
124
+ q_fp8 = torch.empty((n_rows, 128), dtype=torch.float8_e4m3fn, device=q.device)
125
+ q_scale = torch.empty((n_rows, 1), dtype=torch.float32, device=q.device)
126
+ if n_rows == 0:
127
+ return q_fp8, q_scale
128
+ BLOCK_R = 32
129
+ grid = (triton.cdiv(n_rows, BLOCK_R),)
130
+ _fwht_quant_kernel[grid](q, q_fp8, q_scale, n_rows, BLOCK_R=BLOCK_R, num_warps=2)
131
+ return q_fp8, q_scale
132
+
133
+
134
+# Fused pool compression and cache write.
135
+
136
+
137
+@triton.jit
138
+def _kpool_softmax_rotate_write_cache_kernel(
139
+ buf_fp8_ptr,
140
+ buf_fp32_ptr,
141
+ slot_k_ptr,
142
+ slot_score_ptr,
143
+ ape_ptr,
144
+ loc_ptr,
145
+ write_mask_ptr,
146
+ compressed_k_ptr,
147
+ compressed_scale_ptr,
148
+ slot_k_stride_0,
149
+ slot_k_stride_1,
150
+ slot_score_stride_0,
151
+ slot_score_stride_1,
152
+ ape_stride_0,
153
+ PAGE_SIZE: tl.constexpr,
154
+ BUF_NUMEL_PER_PAGE: tl.constexpr,
155
+ POOL_SIZE: tl.constexpr,
156
+ HEAD_DIM: tl.constexpr,
157
+ S_OFFSET_NBYTES_IN_PAGE: tl.constexpr,
158
+ ROUND_SCALE: tl.constexpr,
159
+ HAS_WRITE_MASK: tl.constexpr,
160
+ RETURN_COMPRESSED: tl.constexpr,
161
+ WRITE_CACHE: tl.constexpr,
162
+ BLOCK_D: tl.constexpr,
163
+):
164
+ """One program per pool. softmax(slot_score+ape)-weighted sum of slot_k ->
165
+ Hadamard-128 -> per-vector fp8 absmax quant -> write to cache at ``loc``."""
166
+ row = tl.program_id(0)
167
+ do_write = True
168
+ if HAS_WRITE_MASK:
169
+ do_write = tl.load(write_mask_ptr + row)
170
+
171
+ offs = tl.arange(0, BLOCK_D)
172
+ mask = (offs < HEAD_DIM) & do_write
173
+
174
+ # --- Pass 1: per-dim max over the pool (softmax numerical stability) ---
175
+ max_score = tl.full((BLOCK_D,), -float("inf"), tl.float32)
176
+ for slot in tl.static_range(0, POOL_SIZE):
177
+ score = tl.load(
178
+ slot_score_ptr
179
+ + row * slot_score_stride_0
180
+ + slot * slot_score_stride_1
181
+ + offs,
182
+ mask=mask,
183
+ other=0.0,
184
+ ).to(tl.float32)
185
+ score += tl.load(ape_ptr + slot * ape_stride_0 + offs, mask=mask, other=0.0).to(
186
+ tl.float32
187
+ )
188
+ max_score = tl.maximum(max_score, score)
189
+
190
+ # --- Pass 2: softmax-weighted sum of K ---
191
+ acc = tl.full((BLOCK_D,), 0.0, tl.float32)
192
+ denom = tl.full((BLOCK_D,), 0.0, tl.float32)
193
+ for slot in tl.static_range(0, POOL_SIZE):
194
+ score = tl.load(
195
+ slot_score_ptr
196
+ + row * slot_score_stride_0
197
+ + slot * slot_score_stride_1
198
+ + offs,
199
+ mask=mask,
200
+ other=0.0,
201
+ ).to(tl.float32)
202
+ score += tl.load(ape_ptr + slot * ape_stride_0 + offs, mask=mask, other=0.0).to(
203
+ tl.float32
204
+ )
205
+ prob = tl.exp(score - max_score)
206
+ denom += prob
207
+ k = tl.load(
208
+ slot_k_ptr + row * slot_k_stride_0 + slot * slot_k_stride_1 + offs,
209
+ mask=mask,
210
+ other=0.0,
211
+ ).to(tl.float32)
212
+ acc += k * prob
213
+
214
+ x = acc / denom
215
+ x = tl.where(do_write, x, 0.0).to(tl.bfloat16).to(tl.float32)
216
+
217
+ # Match the unfused pooled-K path's bf16 precision before quantization.
218
+ x = _hadamard128(x).to(tl.bfloat16).to(tl.float32)
219
+
220
+ # --- per-vector absmax fp8 quant ---
221
+ fp8_max = 448.0
222
+ fp8_max_inv = 1.0 / fp8_max
223
+ absmax = tl.max(tl.abs(x), axis=0)
224
+ absmax = tl.maximum(absmax, 1e-4)
225
+ if ROUND_SCALE:
226
+ scale = tl.exp2(tl.ceil(tl.log2(absmax * fp8_max_inv)))
227
+ else:
228
+ scale = absmax * fp8_max_inv
229
+ quantized = x / scale
230
+ quantized = tl.minimum(tl.maximum(quantized, -fp8_max), fp8_max)
231
+
232
+ if WRITE_CACHE:
233
+ loc = tl.load(loc_ptr + row, mask=do_write, other=0)
234
+ loc_page_index = loc // PAGE_SIZE
235
+ loc_token_offset_in_page = loc % PAGE_SIZE
236
+ out_k_offsets = (
237
+ loc_page_index * BUF_NUMEL_PER_PAGE
238
+ + loc_token_offset_in_page * HEAD_DIM
239
+ + offs
240
+ )
241
+ out_s_offset = (
242
+ loc_page_index * BUF_NUMEL_PER_PAGE // 4
243
+ + S_OFFSET_NBYTES_IN_PAGE // 4
244
+ + loc_token_offset_in_page
245
+ )
246
+ tl.store(buf_fp8_ptr + out_k_offsets, quantized, mask=mask)
247
+ tl.store(buf_fp32_ptr + out_s_offset, scale, mask=do_write)
248
+
249
+ if RETURN_COMPRESSED:
250
+ tl.store(
251
+ compressed_k_ptr + row * HEAD_DIM + offs,
252
+ quantized,
253
+ mask=offs < HEAD_DIM,
254
+ )
255
+ tl.store(compressed_scale_ptr + row, scale)
256
+
257
+
258
+def kpool_compress_and_write_cache(
259
+ kv_cache: torch.Tensor,
260
+ slot_k: torch.Tensor,
261
+ slot_score: torch.Tensor,
262
+ ape: torch.Tensor,
263
+ loc: torch.Tensor,
264
+ pool_size: int,
265
+ head_dim: int = INDEX_HEAD_DIM,
266
+ write_mask: torch.Tensor | None = None,
267
+ round_scale: bool = True,
268
+ return_compressed: bool = False,
269
+ write_cache: bool = True,
270
+):
271
+ """Compress ``pool_size`` tokens into one fp8 K and write at ``loc``.
272
+
273
+ Args:
274
+ kv_cache: indexer K cache ``[num_blocks, block_size, head_dim+4]`` uint8.
275
+ slot_k: ``[n_pools, pool_size, head_dim]`` bf16 — raw per-token K.
276
+ slot_score: ``[n_pools, pool_size, head_dim]`` — per-token gate score.
277
+ ape: ``[pool_size, head_dim]`` fp32 — per-slot position bias.
278
+ loc: ``[n_pools]`` int64 — flat physical slot per pool.
279
+ """
280
+ assert slot_k.ndim == 3
281
+ assert slot_score.shape == slot_k.shape
282
+ assert ape.shape == slot_k.shape[1:]
283
+ assert slot_k.shape[2] == head_dim
284
+ assert slot_k.dtype == torch.bfloat16
285
+ assert ape.dtype == torch.float32
286
+ assert kv_cache.dtype == torch.uint8
287
+ assert loc.dtype == torch.int64
288
+ assert write_cache or return_compressed
289
+
290
+ page_size = kv_cache.shape[1]
291
+ buf = kv_cache
292
+ slot_k = slot_k.contiguous()
293
+ slot_score = slot_score.contiguous()
294
+ ape = ape.contiguous()
295
+ loc = loc.contiguous()
296
+ if write_mask is None:
297
+ write_mask = torch.empty((1,), dtype=torch.bool, device=slot_k.device)
298
+ has_write_mask = False
299
+ else:
300
+ assert write_mask.shape == (slot_k.shape[0],)
301
+ write_mask = write_mask.contiguous()
302
+ has_write_mask = True
303
+ assert not return_compressed
304
+
305
+ if slot_k.shape[0] == 0:
306
+ if return_compressed:
307
+ return (
308
+ torch.empty(
309
+ (0, head_dim),
310
+ dtype=torch.float8_e4m3fn,
311
+ device=slot_k.device,
312
+ ),
313
+ torch.empty((0,), dtype=torch.float32, device=slot_k.device),
314
+ )
315
+ return None
316
+
317
+ buf_fp8 = buf.view(torch.float8_e4m3fn)
318
+ buf_fp32 = buf.view(torch.float32)
319
+ # bytes per page (last dim of kv_cache) viewed as uint8
320
+ buf_numel_per_page = buf.stride(0)
321
+ s_offset_nbytes_in_page = page_size * head_dim
322
+
323
+ if return_compressed:
324
+ compressed_k = torch.empty(
325
+ (slot_k.shape[0], head_dim),
326
+ dtype=torch.float8_e4m3fn,
327
+ device=slot_k.device,
328
+ )
329
+ compressed_scale = torch.empty(
330
+ (slot_k.shape[0],), dtype=torch.float32, device=slot_k.device
331
+ )
332
+ else:
333
+ compressed_k = buf_fp8
334
+ compressed_scale = buf_fp32
335
+
336
+ _kpool_softmax_rotate_write_cache_kernel[(slot_k.shape[0],)](
337
+ buf_fp8,
338
+ buf_fp32,
339
+ slot_k,
340
+ slot_score,
341
+ ape,
342
+ loc,
343
+ write_mask,
344
+ compressed_k,
345
+ compressed_scale,
346
+ slot_k.stride(0),
347
+ slot_k.stride(1),
348
+ slot_score.stride(0),
349
+ slot_score.stride(1),
350
+ ape.stride(0),
351
+ PAGE_SIZE=page_size,
352
+ BUF_NUMEL_PER_PAGE=buf_numel_per_page,
353
+ POOL_SIZE=slot_k.shape[1],
354
+ HEAD_DIM=head_dim,
355
+ S_OFFSET_NBYTES_IN_PAGE=s_offset_nbytes_in_page,
356
+ ROUND_SCALE=round_scale,
357
+ HAS_WRITE_MASK=has_write_mask,
358
+ RETURN_COMPRESSED=return_compressed,
359
+ WRITE_CACHE=write_cache,
360
+ BLOCK_D=triton.next_power_of_2(head_dim),
361
+ )
362
+
363
+ if return_compressed:
364
+ return compressed_k, compressed_scale
365
+ return None
366
+
367
+
368
+# Seed each request's incomplete pool into its paged tail during prefill.
369
+
370
+
371
+@triton.jit
372
+def _kpool_tail_seed_kernel(
373
+ key_ptr,
374
+ score_ptr,
375
+ tslot_ptr,
376
+ tail_ptr,
377
+ n_tokens,
378
+ HEAD_DIM: tl.constexpr,
379
+ KPOOL: tl.constexpr,
380
+ BLOCK_D: tl.constexpr,
381
+):
382
+ """Copy token ``i``'s raw K + gate into its request's tail block.
383
+
384
+ Token ``i`` is among its request's last KPOOL tokens iff the token KPOOL
385
+ ahead belongs to a different tail block (or is past the batch / padding,
386
+ slot < 0). ``tslot = block * KPOOL + pos % KPOOL``; the destination is
387
+ ``tail[block, {0:K, 1:score}, pos % KPOOL, :]``.
388
+ """
389
+ i = tl.program_id(0)
390
+ t = tl.load(tslot_ptr + i).to(tl.int64)
391
+ if t < 0:
392
+ return
393
+ blk = t // KPOOL # t >= 0 here, so trunc == floor
394
+ ahead = tl.load(tslot_ptr + i + KPOOL, mask=i + KPOOL < n_tokens, other=-1).to(
395
+ tl.int64
396
+ )
397
+ # Match the torch semantics exactly: a negative ahead slot floors to a
398
+ # block id that differs from every real block -> token is in the tail.
399
+ # Only divide non-negative slots (Triton int div truncates, torch floors).
400
+ if ahead >= 0 and ahead // KPOOL == blk:
401
+ return
402
+ offs = tl.arange(0, BLOCK_D)
403
+ m = offs < HEAD_DIM
404
+ base = (blk * 2 * KPOOL + t % KPOOL) * HEAD_DIM
405
+ k = tl.load(key_ptr + i * HEAD_DIM + offs, mask=m)
406
+ s = tl.load(score_ptr + i * HEAD_DIM + offs, mask=m)
407
+ tl.store(tail_ptr + base + offs, k, mask=m)
408
+ tl.store(tail_ptr + base + KPOOL * HEAD_DIM + offs, s, mask=m)
409
+
410
+
411
+def kpool_seed_tail_cache(
412
+ tail_kv_cache: torch.Tensor,
413
+ key: torch.Tensor,
414
+ gate_score: torch.Tensor,
415
+ tslot: torch.Tensor,
416
+ kpool: int,
417
+ head_dim: int = INDEX_HEAD_DIM,
418
+) -> None:
419
+ """Seed the paged tail cache from a prefill batch (see the kernel)."""
420
+ assert tail_kv_cache.dtype == torch.bfloat16
421
+ assert key.dtype == torch.bfloat16
422
+ n = tslot.shape[0]
423
+ if n == 0:
424
+ return
425
+ _kpool_tail_seed_kernel[(n,)](
426
+ key,
427
+ gate_score,
428
+ tslot,
429
+ tail_kv_cache,
430
+ n,
431
+ HEAD_DIM=head_dim,
432
+ KPOOL=kpool,
433
+ BLOCK_D=triton.next_power_of_2(head_dim),
434
+ )
435
+
436
+
437
+# Update each request's tail during decode and write completed pools.
438
+
439
+
440
+@triton.jit
441
+def _kpool_decode_update_batched_kernel(
442
+ buf_fp8_ptr,
443
+ buf_fp32_ptr,
444
+ tail_kv_ptr,
445
+ tail_slot_mapping_ptr, # [B, NEXT_N] int32
446
+ key_ptr, # [B, NEXT_N, HEAD_DIM] bf16
447
+ key_stride_b,
448
+ key_stride_t,
449
+ slot_score_ptr, # [B, NEXT_N, HEAD_DIM] bf16
450
+ ss_stride_b,
451
+ ss_stride_t,
452
+ ape_ptr,
453
+ ape_stride_0,
454
+ slot_mapping_ptr, # [B, NEXT_N] int32
455
+ positions_ptr, # [B, NEXT_N] int32
456
+ NEXT_N, # runtime token count per request (no .item() needed)
457
+ PAGE_SIZE: tl.constexpr,
458
+ BUF_NUMEL_PER_PAGE: tl.constexpr,
459
+ POOL_SIZE: tl.constexpr,
460
+ TAIL_BLOCK_ELEMS: tl.constexpr,
461
+ KPOOL_HEAD: tl.constexpr,
462
+ HEAD_DIM: tl.constexpr,
463
+ S_OFFSET_NBYTES_IN_PAGE: tl.constexpr,
464
+ ROUND_SCALE: tl.constexpr,
465
+ BLOCK_D: tl.constexpr,
466
+):
467
+ """One program per request; iterates its NEXT_N verify tokens in order.
468
+
469
+ Replaces the caller's per-token sequential launch loop. The intra-request
470
+ iteration MUST stay in position order: a pool-completion at token t* reads
471
+ the tail-ring slots that tokens t < t* (same request) just stashed in this
472
+ same invocation. ``tl.range`` iterates sequentially within the program, so
473
+ those stashes are visible to the later completion read. Cross-request
474
+ programs are independent (distinct tail blocks). With NEXT_N < POOL_SIZE
475
+ (the spec-verify case: NEXT_N ~= num_spec+1, POOL_SIZE=16) at most one
476
+ completion can occur per request per call, but the ordered loop is correct
477
+ for any NEXT_N.
478
+ """
479
+ req = tl.program_id(0)
480
+ offs = tl.arange(0, BLOCK_D)
481
+ dim_mask = offs < HEAD_DIM
482
+
483
+ for t in tl.range(0, NEXT_N):
484
+ idx = req * NEXT_N + t
485
+ cache_loc = tl.load(slot_mapping_ptr + idx)
486
+ pos = tl.load(positions_ptr + idx)
487
+ safe_pos = tl.maximum(pos, 0)
488
+ pos_valid = (cache_loc >= 0) & (pos >= 0)
489
+
490
+ slot = safe_pos % POOL_SIZE
491
+ phys_slot = safe_pos % POOL_SIZE
492
+
493
+ # Derive the tail block from THIS token's tail_slot (the request's block
494
+ # is constant across a pool, but a padded / invalid entry carries a
495
+ # negative sentinel -- reading it from token 0 would poison every
496
+ # token's base address). Clamp so an invalid entry can never form an
497
+ # out-of-bounds base; the accesses below are gated on pos_valid anyway.
498
+ tail_slot = tl.load(tail_slot_mapping_ptr + idx)
499
+ block = tl.maximum(tail_slot, 0).to(tl.int64) // POOL_SIZE
500
+ block_base = block * TAIL_BLOCK_ELEMS
501
+
502
+ # The tail-ring stash must run for EVERY real token, so it is gated on
503
+ # the token-granular tail slot -- not on `pos_valid`, which keys off the
504
+ # POOL-granular `slot_mapping` and is therefore only true on the pool's
505
+ # last token. Gating the stash on pos_valid dropped every intra-pool
506
+ # token, so a decode-built pool compressed 3 stale ring entries (the
507
+ # prefill-seeded prompt tail, frozen forever) plus the current token.
508
+ stash_valid = (pos >= 0) & (tail_slot >= 0)
509
+
510
+ key = tl.load(
511
+ key_ptr + req * key_stride_b + t * key_stride_t + offs,
512
+ mask=dim_mask,
513
+ other=0.0,
514
+ ).to(tl.float32)
515
+ score_current = tl.load(
516
+ slot_score_ptr + req * ss_stride_b + t * ss_stride_t + offs,
517
+ mask=dim_mask,
518
+ other=0.0,
519
+ ).to(tl.float32)
520
+
521
+ if pos_valid & (slot == POOL_SIZE - 1):
522
+ pool_logical_start = safe_pos - slot
523
+
524
+ max_score = tl.full((BLOCK_D,), -float("inf"), tl.float32)
525
+ for pool_slot in tl.static_range(0, POOL_SIZE):
526
+ is_current = pool_slot == slot
527
+ phys = (pool_logical_start + pool_slot) % POOL_SIZE
528
+ score_buf = tl.load(
529
+ tail_kv_ptr + block_base + KPOOL_HEAD + phys * HEAD_DIM + offs,
530
+ mask=dim_mask,
531
+ other=0.0,
532
+ ).to(tl.float32)
533
+ score = tl.where(is_current, score_current, score_buf)
534
+ score += tl.load(
535
+ ape_ptr + pool_slot * ape_stride_0 + offs,
536
+ mask=dim_mask,
537
+ other=0.0,
538
+ ).to(tl.float32)
539
+ max_score = tl.maximum(max_score, score)
540
+
541
+ acc = tl.full((BLOCK_D,), 0.0, tl.float32)
542
+ denom = tl.full((BLOCK_D,), 0.0, tl.float32)
543
+ for pool_slot in tl.static_range(0, POOL_SIZE):
544
+ is_current = pool_slot == slot
545
+ phys = (pool_logical_start + pool_slot) % POOL_SIZE
546
+ score_buf = tl.load(
547
+ tail_kv_ptr + block_base + KPOOL_HEAD + phys * HEAD_DIM + offs,
548
+ mask=dim_mask,
549
+ other=0.0,
550
+ ).to(tl.float32)
551
+ score = tl.where(is_current, score_current, score_buf)
552
+ score += tl.load(
553
+ ape_ptr + pool_slot * ape_stride_0 + offs,
554
+ mask=dim_mask,
555
+ other=0.0,
556
+ ).to(tl.float32)
557
+ prob = tl.exp(score - max_score)
558
+ denom += prob
559
+ k_buf = tl.load(
560
+ tail_kv_ptr + block_base + phys * HEAD_DIM + offs,
561
+ mask=dim_mask,
562
+ other=0.0,
563
+ ).to(tl.float32)
564
+ k = tl.where(is_current, key, k_buf)
565
+ acc += k * prob
566
+
567
+ x = (acc / denom).to(tl.bfloat16).to(tl.float32)
568
+ x = _hadamard128(x).to(tl.bfloat16).to(tl.float32)
569
+
570
+ fp8_max = 448.0
571
+ fp8_max_inv = 1.0 / fp8_max
572
+ absmax = tl.maximum(tl.max(tl.abs(x), axis=0), 1e-4)
573
+ if ROUND_SCALE:
574
+ scale = tl.exp2(tl.ceil(tl.log2(absmax * fp8_max_inv)))
575
+ else:
576
+ scale = absmax * fp8_max_inv
577
+ quantized = tl.minimum(tl.maximum(x / scale, -fp8_max), fp8_max)
578
+
579
+ loc = cache_loc.to(tl.int64)
580
+ loc_page_index = loc // PAGE_SIZE
581
+ loc_token_offset_in_page = loc % PAGE_SIZE
582
+ out_k_offsets = (
583
+ loc_page_index * BUF_NUMEL_PER_PAGE
584
+ + loc_token_offset_in_page * HEAD_DIM
585
+ + offs
586
+ )
587
+ out_s_offset = (
588
+ loc_page_index * BUF_NUMEL_PER_PAGE // 4
589
+ + S_OFFSET_NBYTES_IN_PAGE // 4
590
+ + loc_token_offset_in_page
591
+ )
592
+ tl.store(buf_fp8_ptr + out_k_offsets, quantized, mask=dim_mask)
593
+ tl.store(buf_fp32_ptr + out_s_offset, scale)
594
+
595
+ # Stash the current token AFTER any completion read so the completion
596
+ # uses prior stashes (and the current token's own key/score via
597
+ # is_current), then leaves this token for future pools. Order matches
598
+ # the per-token kernel: completion read first, stash second.
599
+ update_mask = dim_mask & stash_valid
600
+ tl.store(
601
+ tail_kv_ptr + block_base + phys_slot * HEAD_DIM + offs,
602
+ key,
603
+ mask=update_mask,
604
+ )
605
+ tl.store(
606
+ tail_kv_ptr + block_base + KPOOL_HEAD + phys_slot * HEAD_DIM + offs,
607
+ score_current,
608
+ mask=update_mask,
609
+ )
610
+
611
+
612
+def kpool_decode_update_and_maybe_write_cache_batched(
613
+ kv_cache: torch.Tensor,
614
+ tail_kv_cache: torch.Tensor,
615
+ tail_slot_mapping: torch.Tensor,
616
+ key: torch.Tensor,
617
+ slot_score: torch.Tensor,
618
+ ape: torch.Tensor,
619
+ slot_mapping: torch.Tensor,
620
+ positions: torch.Tensor,
621
+ pool_size: int,
622
+ head_dim: int = INDEX_HEAD_DIM,
623
+ round_scale: bool = True,
624
+) -> None:
625
+ """Batched decode-step kpool update for spec verify (``next_n > 1``).
626
+
627
+ One launch replaces the caller's per-token loop. Inputs are grouped per
628
+ request: ``[num_requests, next_n, ...]``. Each program handles one
629
+ request's ``next_n`` tokens in position order (see the kernel docstring for
630
+ why ordering is required for pool-completion correctness).
631
+
632
+ Plain decode (``next_n == 1``) is handled here too — the kernel collapses
633
+ to a single-iteration loop.
634
+
635
+ Args:
636
+ kv_cache: indexer K cache ``[num_blocks, block_size, head_dim+4]`` uint8.
637
+ tail_kv_cache: paged tail cache ``[num_blocks, 2, pool_size, head_dim]``
638
+ bf16 (K at half 0, gate score at half 1).
639
+ tail_slot_mapping: ``[num_requests, next_n]`` int32.
640
+ key / slot_score: ``[num_requests, next_n, head_dim]`` bf16.
641
+ ape: ``[pool_size, head_dim]`` fp32.
642
+ slot_mapping / positions: ``[num_requests, next_n]`` int32.
643
+ """
644
+ num_requests, next_n = key.shape[0], key.shape[1]
645
+ if num_requests == 0 or next_n == 0:
646
+ return
647
+ assert tail_kv_cache.ndim == 4
648
+ assert tail_kv_cache.shape[1] == 2
649
+ assert tail_kv_cache.shape[2] == pool_size
650
+ assert tail_kv_cache.shape[3] == head_dim
651
+ assert tail_kv_cache.dtype == torch.bfloat16
652
+ assert key.ndim == 3 and key.shape[2] == head_dim
653
+ assert slot_score.shape == key.shape
654
+ assert ape.shape == (pool_size, head_dim)
655
+ assert tail_slot_mapping.shape == (num_requests, next_n)
656
+ assert slot_mapping.shape == (num_requests, next_n)
657
+ assert positions.shape == (num_requests, next_n)
658
+ assert key.dtype == torch.bfloat16
659
+ assert slot_score.dtype == torch.bfloat16
660
+ assert ape.dtype == torch.float32
661
+ assert kv_cache.dtype == torch.uint8
662
+
663
+ page_size = kv_cache.shape[1]
664
+ buf = kv_cache
665
+ buf_fp8 = buf.view(torch.float8_e4m3fn)
666
+ buf_fp32 = buf.view(torch.float32)
667
+
668
+ # The kernel indexes the int tensors as ``req * next_n + t`` (row-major),
669
+ # so they must be contiguous. Callers pass either a view of a contiguous
670
+ # slice or a freshly scattered tensor, making these no-ops; the calls guard
671
+ # against a future caller handing over a strided view.
672
+ tail_slot_mapping = tail_slot_mapping.contiguous()
673
+ slot_mapping = slot_mapping.contiguous()
674
+ positions = positions.contiguous()
675
+
676
+ _kpool_decode_update_batched_kernel[(num_requests,)](
677
+ buf_fp8,
678
+ buf_fp32,
679
+ tail_kv_cache,
680
+ tail_slot_mapping,
681
+ key,
682
+ key.stride(0),
683
+ key.stride(1),
684
+ slot_score,
685
+ slot_score.stride(0),
686
+ slot_score.stride(1),
687
+ ape,
688
+ ape.stride(0),
689
+ slot_mapping,
690
+ positions,
691
+ next_n,
692
+ PAGE_SIZE=page_size,
693
+ BUF_NUMEL_PER_PAGE=buf.stride(0),
694
+ POOL_SIZE=pool_size,
695
+ TAIL_BLOCK_ELEMS=tail_kv_cache.stride(0),
696
+ KPOOL_HEAD=tail_kv_cache.stride(1),
697
+ HEAD_DIM=head_dim,
698
+ S_OFFSET_NBYTES_IN_PAGE=page_size * head_dim,
699
+ ROUND_SCALE=round_scale,
700
+ BLOCK_D=triton.next_power_of_2(head_dim),
701
+ )
702
+
703
+
704
+# Pool-level top-k helpers.
705
+
706
+
707
+def history_group_budget_for_topk(topk: int, pool_size: int) -> int:
708
+ """Number of pools to select so that expanding yields ``topk`` tokens."""
709
+ assert topk % pool_size == 0
710
+ return topk // pool_size
711
+
712
+
713
+def expand_pools_to_tokens(
714
+ group_ids: torch.Tensor,
715
+ group_valid: torch.Tensor,
716
+ topk: int,
717
+ pool_size: int,
718
+ page_table: torch.Tensor | None = None,
719
+ topk_offsets: torch.Tensor | None = None,
720
+) -> torch.Tensor:
721
+ """Expand selected full-pool ids to a strict-width token topk tensor."""
722
+ assert group_ids.ndim == 2
723
+ assert group_valid.shape == group_ids.shape
724
+ assert topk % pool_size == 0
725
+ assert group_ids.shape[1] == history_group_budget_for_topk(topk, pool_size)
726
+ assert page_table is None or topk_offsets is None
727
+
728
+ device = group_ids.device
729
+ offsets = torch.arange(pool_size, device=device, dtype=torch.int64)
730
+ token_ids = group_ids.to(torch.int64).unsqueeze(-1) * pool_size + offsets
731
+ token_ids = token_ids.reshape(group_ids.shape[0], topk)
732
+ valid = (
733
+ group_valid.unsqueeze(-1)
734
+ .expand(-1, -1, pool_size)
735
+ .reshape(group_ids.shape[0], topk)
736
+ )
737
+
738
+ if page_table is not None:
739
+ assert page_table.ndim == 2
740
+ safe_ids = token_ids.clamp(min=0, max=page_table.shape[1] - 1)
741
+ output = torch.gather(page_table, dim=1, index=safe_ids).to(torch.int32)
742
+ elif topk_offsets is not None:
743
+ if topk_offsets.ndim == 2:
744
+ assert topk_offsets.shape[1] == 1
745
+ topk_offsets = topk_offsets.squeeze(1)
746
+ output = (token_ids + topk_offsets.to(torch.int64).unsqueeze(1)).to(torch.int32)
747
+ else:
748
+ output = token_ids.to(torch.int32)
749
+
750
+ return torch.where(valid, output, torch.full_like(output, -1))
751
+
752
+
753
+def append_tail_to_topk(
754
+ topk_result: torch.Tensor,
755
+ seq_lens: torch.Tensor,
756
+ pool_lens: torch.Tensor,
757
+ pool_size: int,
758
+ page_table: torch.Tensor | None = None,
759
+ topk_offsets: torch.Tensor | None = None,
760
+) -> torch.Tensor:
761
+ """Append non-pooled tail tokens after expanded history tokens.
762
+
763
+ ``index_kpool_always_select_tail`` keeps the (incomplete) trailing pool so
764
+ the most recent tokens are always attended to.
765
+ """
766
+ assert topk_result.dtype == torch.int32
767
+ assert seq_lens.ndim == 1
768
+ assert pool_lens.ndim == 1
769
+
770
+ tail_pool = pool_size - 1
771
+ if tail_pool == 0:
772
+ return topk_result
773
+
774
+ rows, n_cols = topk_result.shape
775
+ out_cols = n_cols + tail_pool
776
+ out = torch.empty(
777
+ (rows, out_cols), dtype=topk_result.dtype, device=topk_result.device
778
+ )
779
+
780
+ # tail tokens: [pool_len*pool_size, seq_len) for each row.
781
+ pool_len = pool_lens.to(torch.int32)
782
+ tail_start = pool_len * pool_size
783
+ seq_len = seq_lens.to(torch.int32)
784
+ tail_count = seq_len - tail_start # in [0, pool_size)
785
+
786
+ cols = torch.arange(out_cols, device=topk_result.device)[None, :]
787
+ history_len = n_cols
788
+ is_history = cols < history_len
789
+ tail_off = cols - history_len
790
+ is_tail = (tail_off >= 0) & (tail_off < tail_count[:, None])
791
+
792
+ # safe_hist must be per-row [rows, out_cols] so the gather reads each row's
793
+ # OWN history. cols is [1, out_cols]; if used directly, gather (which does
794
+ # NOT broadcast the index) would read only row 0 of topk_result, making every
795
+ # query inherit row 0's history (empty for the first token) and lose all its
796
+ # selected tokens — only the per-row tail would survive. This only manifests
797
+ # for multi-row sparse PREFILL (decode has 1 row, so it reads its own row 0).
798
+ safe_hist = torch.minimum(cols, torch.full_like(cols, n_cols - 1)).expand(
799
+ rows, out_cols
800
+ )
801
+ history_val = torch.gather(topk_result, 1, safe_hist)
802
+
803
+ tail_raw = tail_start[:, None] + tail_off
804
+ tail_val = tail_raw.to(torch.int32)
805
+ if page_table is not None:
806
+ safe_tail = tail_raw.clamp(min=0, max=page_table.shape[1] - 1)
807
+ tail_val = torch.gather(page_table, 1, safe_tail).to(torch.int32)
808
+ elif topk_offsets is not None:
809
+ tail_val = (tail_raw + topk_offsets.to(torch.int64).unsqueeze(1)).to(
810
+ torch.int32
811
+ )
812
+
813
+ out = torch.where(is_history, history_val, -1)
814
+ out = torch.where(is_tail, tail_val, out)
815
+ return out
816
+
817
+
818
+@triton.jit
819
+def _expand_pools_and_append_tail_kernel(
820
+ pool_ids_ptr, # [rows, n_groups], int (any int dtype)
821
+ seq_lens_ptr, # [rows], int32 (token-granular seq_len)
822
+ out_ptr, # [rows, out_cols], int32
823
+ topk, # n_groups * pool_size
824
+ out_cols, # topk + pool_size - 1
825
+ POOL_SIZE: tl.constexpr,
826
+ BLOCK_COLS: tl.constexpr,
827
+ pid_s0,
828
+ out_s0,
829
+):
830
+ # Fuses expand_pools_to_tokens + append_tail_to_topk (identity path) into a
831
+ # single kernel. Each program writes one (row, column-tile) of the output.
832
+ row = tl.program_id(0)
833
+ tile = tl.program_id(1)
834
+ cols = tile * BLOCK_COLS + tl.arange(0, BLOCK_COLS)
835
+ mask = cols < out_cols
836
+
837
+ seq_len = tl.load(seq_lens_ptr + row)
838
+ pool_len = seq_len // POOL_SIZE
839
+ tail_start = pool_len * POOL_SIZE
840
+ tail_count = seq_len - tail_start # in [0, POOL_SIZE)
841
+
842
+ # History region [0, topk): expand selected pool g = cols // POOL_SIZE.
843
+ is_history = cols < topk
844
+ g = cols // POOL_SIZE
845
+ o = cols % POOL_SIZE
846
+ pid = tl.load(pool_ids_ptr + row * pid_s0 + g, mask=mask & is_history, other=-1)
847
+ hist_val = (pid * POOL_SIZE + o).to(tl.int32)
848
+ hist_out = tl.where(pid >= 0, hist_val, -1)
849
+
850
+ # Tail region [topk, out_cols): the request's trailing incomplete pool.
851
+ tail_off = cols - topk
852
+ is_tail = (tail_off >= 0) & (tail_off < tail_count)
853
+ tail_val = (tail_start + tail_off).to(tl.int32)
854
+ tail_out = tl.where(is_tail, tail_val, -1)
855
+
856
+ result = tl.where(is_history, hist_out, tail_out)
857
+ tl.store(out_ptr + row * out_s0 + cols, result, mask=mask)
858
+
859
+
860
+def expand_pools_and_append_tail(
861
+ pool_ids: torch.Tensor,
862
+ seq_lens: torch.Tensor,
863
+ pool_size: int,
864
+) -> torch.Tensor:
865
+ """Fuse ``expand_pools_to_tokens`` + ``append_tail_to_topk`` (identity path).
866
+
867
+ Produces the same ``[rows, topk + pool_size - 1]`` int32 output as calling
868
+ the two functions in sequence when neither ``page_table`` nor
869
+ ``topk_offsets`` is passed — the only path used by the GLM-5.3-Flash indexer.
870
+ The kernel derives ``pool_len = seq_len // pool_size`` internally, so the
871
+ caller no longer needs to precompute it. Replaces ~25 elementwise kernels
872
+ with one Triton launch.
873
+ """
874
+ rows, n_groups = pool_ids.shape
875
+ topk = n_groups * pool_size
876
+ out_cols = topk + pool_size - 1
877
+ out = torch.empty((rows, out_cols), dtype=torch.int32, device=pool_ids.device)
878
+ BLOCK_COLS = 128
879
+ n_tiles = triton.cdiv(out_cols, BLOCK_COLS)
880
+ _expand_pools_and_append_tail_kernel[(rows, n_tiles)](
881
+ pool_ids,
882
+ seq_lens,
883
+ out,
884
+ topk,
885
+ out_cols,
886
+ POOL_SIZE=pool_size,
887
+ BLOCK_COLS=BLOCK_COLS,
888
+ pid_s0=pool_ids.stride(0),
889
+ out_s0=out.stride(0),
890
+ )
891
+ return out
vllm/models/glm5next/nvidia/ops/third_party/__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/glm5next/nvidia/ops/third_party/kda/__init__.py
new
+6
@@ -0,0 +1,6 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from .kernels import chunk_kda_with_fused_gate, fused_recurrent_kda
5
+
6
+__all__ = ["chunk_kda_with_fused_gate", "fused_recurrent_kda"]
vllm/models/glm5next/nvidia/ops/third_party/kda/fused_recurrent.py
new
+656
@@ -0,0 +1,656 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang
4
+# mypy: ignore-errors
5
+#
6
+# This file contains code copied from the flash-linear-attention project.
7
+# The original source code was licensed under the MIT license and included
8
+# the following copyright notice:
9
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
10
+# ruff: noqa: E501
11
+
12
+import torch
13
+
14
+from vllm.third_party.flash_linear_attention.ops.op import exp
15
+from vllm.triton_utils import tl, triton
16
+
17
+
18
+@triton.heuristics(
19
+ {
20
+ "USE_INITIAL_STATE": lambda args: args["h0"] is not None,
21
+ "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
22
+ "IS_CONTINUOUS_BATCHING": lambda args: args["ssm_state_indices"] is not None,
23
+ "IS_SPEC_DECODING": lambda args: args["num_accepted_tokens"] is not None,
24
+ }
25
+)
26
+@triton.jit(do_not_specialize=["N", "T"])
27
+def fused_recurrent_gated_delta_rule_fwd_kernel(
28
+ q,
29
+ k,
30
+ v,
31
+ g,
32
+ beta,
33
+ o,
34
+ h0,
35
+ ht,
36
+ cu_seqlens,
37
+ ssm_state_indices,
38
+ num_accepted_tokens,
39
+ a_log,
40
+ g_bias,
41
+ scale,
42
+ N: tl.int64, # num of sequences
43
+ T: tl.int64, # num of tokens
44
+ B: tl.constexpr,
45
+ H: tl.constexpr,
46
+ HV: tl.constexpr,
47
+ K: tl.constexpr,
48
+ V: tl.constexpr,
49
+ BK: tl.constexpr,
50
+ BV: tl.constexpr,
51
+ stride_init_state_token: tl.constexpr,
52
+ stride_final_state_token: tl.constexpr,
53
+ stride_indices_seq: tl.constexpr,
54
+ stride_indices_tok: tl.constexpr,
55
+ USE_INITIAL_STATE: tl.constexpr, # whether to use initial state
56
+ INPLACE_FINAL_STATE: tl.constexpr, # whether to store final state inplace
57
+ IS_BETA_HEADWISE: tl.constexpr, # whether beta is headwise vector or scalar,
58
+ USE_QK_L2NORM_IN_KERNEL: tl.constexpr,
59
+ IS_VARLEN: tl.constexpr,
60
+ IS_CONTINUOUS_BATCHING: tl.constexpr,
61
+ IS_SPEC_DECODING: tl.constexpr,
62
+ IS_KDA: tl.constexpr,
63
+ SIGMOID_BETA: tl.constexpr, # beta holds raw logits; sigmoid at fp32 load
64
+ COMPUTE_GATE: tl.constexpr, # g holds raw logits; KDA gate computed in-kernel
65
+ SAFE_GATE: tl.constexpr, # bounded gate variant (only branch implemented)
66
+ LOWER_BOUND: tl.constexpr,
67
+):
68
+ i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
69
+ i_n, i_hv = i_nh // HV, i_nh % HV
70
+ i_h = i_hv // (HV // H)
71
+ if IS_VARLEN:
72
+ bos, eos = (
73
+ tl.load(cu_seqlens + i_n).to(tl.int64),
74
+ tl.load(cu_seqlens + i_n + 1).to(tl.int64),
75
+ )
76
+ all = T
77
+ T = eos - bos
78
+ else:
79
+ bos, eos = i_n * T, i_n * T + T
80
+ all = B * T
81
+
82
+ if T == 0:
83
+ # no tokens to process for this sequence
84
+ return
85
+
86
+ o_k = i_k * BK + tl.arange(0, BK)
87
+ o_v = i_v * BV + tl.arange(0, BV)
88
+
89
+ p_q = q + (bos * H + i_h) * K + o_k
90
+ p_k = k + (bos * H + i_h) * K + o_k
91
+ p_v = v + (bos * HV + i_hv) * V + o_v
92
+ if IS_BETA_HEADWISE:
93
+ p_beta = beta + (bos * HV + i_hv) * V + o_v
94
+ else:
95
+ p_beta = beta + bos * HV + i_hv
96
+
97
+ if not IS_KDA:
98
+ p_g = g + bos * HV + i_hv
99
+ else:
100
+ p_gk = g + (bos * HV + i_hv) * K + o_k
101
+
102
+ # Per-head gate amplitude, hoisted out of the token loop (COMPUTE_GATE).
103
+ if COMPUTE_GATE:
104
+ b_a_log = tl.exp(tl.load(a_log + i_h).to(tl.float32))
105
+
106
+ p_o = o + ((i_k * all + bos) * HV + i_hv) * V + o_v
107
+
108
+ mask_k = o_k < K
109
+ mask_v = o_v < V
110
+ mask_h = mask_v[:, None] & mask_k[None, :]
111
+
112
+ b_h = tl.zeros([BV, BK], dtype=tl.float32)
113
+ if USE_INITIAL_STATE:
114
+ if IS_CONTINUOUS_BATCHING:
115
+ if IS_SPEC_DECODING:
116
+ i_t = tl.load(num_accepted_tokens + i_n).to(tl.int64) - 1
117
+ else:
118
+ i_t = 0
119
+ # Load state index and check for invalid entries
120
+ state_idx = tl.load(ssm_state_indices + i_n * stride_indices_seq + i_t).to(
121
+ tl.int64
122
+ )
123
+ # Skip if state index is invalid (NULL_BLOCK_ID=0)
124
+ if state_idx <= 0:
125
+ return
126
+ p_h0 = h0 + state_idx * stride_init_state_token
127
+ else:
128
+ p_h0 = h0 + bos * HV * V * K
129
+ p_h0 = p_h0 + i_hv * V * K + o_v[:, None] * K + o_k[None, :]
130
+ b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32)
131
+
132
+ for i_t in range(0, T):
133
+ b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32)
134
+ b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32)
135
+ b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32)
136
+
137
+ if USE_QK_L2NORM_IN_KERNEL:
138
+ b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6)
139
+ b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6)
140
+ b_q = b_q * scale
141
+ # [BV, BK]
142
+ if not IS_KDA:
143
+ b_g = tl.load(p_g).to(tl.float32)
144
+ b_h *= exp(b_g)
145
+ else:
146
+ b_gk = tl.load(p_gk).to(tl.float32)
147
+ if COMPUTE_GATE:
148
+ # Replicates kda_gate_fwd_kernel's SAFE_GATE branch
149
+ # bit-for-bit (same tl.exp, same fp32 math; the intermediate
150
+ # gate value this replaces was stored/reloaded as fp32,
151
+ # which is lossless): y = lb / (1 + exp(-exp(A)*(g+bias))).
152
+ b_gk += tl.load(g_bias + i_h * K + o_k, mask=mask_k, other=0.0).to(
153
+ tl.float32
154
+ )
155
+ b_gk = LOWER_BOUND / (1.0 + tl.exp(-(b_a_log * b_gk)))
156
+ b_h *= exp(b_gk[None, :])
157
+ # [BV]
158
+ b_v -= tl.sum(b_h * b_k[None, :], 1)
159
+ if IS_BETA_HEADWISE:
160
+ b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32)
161
+ else:
162
+ b_beta = tl.load(p_beta).to(tl.float32)
163
+ # Matches torch's `x.float().sigmoid()` pre-computation bit-for-bit
164
+ # on the input side (bf16->fp32 is exact); only the sigmoid impl itself
165
+ # can differ by <=1 ULP.
166
+ if SIGMOID_BETA:
167
+ b_beta = tl.sigmoid(b_beta)
168
+ b_v *= b_beta
169
+ # [BV, BK]
170
+ b_h += b_v[:, None] * b_k[None, :]
171
+ # [BV]
172
+ b_o = tl.sum(b_h * b_q[None, :], 1)
173
+ tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v)
174
+
175
+ # keep the states for multi-query tokens
176
+ if INPLACE_FINAL_STATE:
177
+ # Load state index and check for invalid entries
178
+ final_state_idx = tl.load(
179
+ ssm_state_indices + i_n * stride_indices_seq + i_t
180
+ ).to(tl.int64)
181
+ # Only store if state index is valid (not NULL_BLOCK_ID=0)
182
+ if final_state_idx > 0:
183
+ p_ht = ht + final_state_idx * stride_final_state_token
184
+ p_ht = p_ht + i_hv * V * K + o_v[:, None] * K + o_k[None, :]
185
+ tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h)
186
+ else:
187
+ p_ht = ht + (bos + i_t) * stride_final_state_token
188
+ p_ht = p_ht + i_hv * V * K + o_v[:, None] * K + o_k[None, :]
189
+ tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h)
190
+
191
+ p_q += H * K
192
+ p_k += H * K
193
+ p_o += HV * V
194
+ p_v += HV * V
195
+ if not IS_KDA:
196
+ p_g += HV
197
+ else:
198
+ p_gk += HV * K
199
+ p_beta += HV * (V if IS_BETA_HEADWISE else 1)
200
+
201
+
202
+def fused_recurrent_gated_delta_rule_fwd(
203
+ q: torch.Tensor,
204
+ k: torch.Tensor,
205
+ v: torch.Tensor,
206
+ g: torch.Tensor,
207
+ beta: torch.Tensor,
208
+ scale: float,
209
+ initial_state: torch.Tensor,
210
+ inplace_final_state: bool = True,
211
+ cu_seqlens: torch.Tensor | None = None,
212
+ ssm_state_indices: torch.Tensor | None = None,
213
+ num_accepted_tokens: torch.Tensor | None = None,
214
+ use_qk_l2norm_in_kernel: bool = False,
215
+) -> tuple[torch.Tensor, torch.Tensor]:
216
+ B, T, H, K, V = *k.shape, v.shape[-1]
217
+ HV = v.shape[2]
218
+ N = B if cu_seqlens is None else len(cu_seqlens) - 1
219
+ BK, BV = triton.next_power_of_2(K), min(triton.next_power_of_2(V), 32)
220
+ NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV)
221
+ assert NK == 1, "NK > 1 is not supported yet"
222
+ num_stages = 3
223
+ num_warps = 1
224
+
225
+ o = q.new_empty(NK, *v.shape)
226
+ if inplace_final_state:
227
+ final_state = initial_state
228
+ else:
229
+ final_state = q.new_empty(T, HV, V, K, dtype=initial_state.dtype)
230
+
231
+ stride_init_state_token = initial_state.stride(0)
232
+ stride_final_state_token = final_state.stride(0)
233
+
234
+ if ssm_state_indices is None:
235
+ stride_indices_seq, stride_indices_tok = 1, 1
236
+ elif ssm_state_indices.ndim == 1:
237
+ stride_indices_seq, stride_indices_tok = ssm_state_indices.stride(0), 1
238
+ else:
239
+ stride_indices_seq, stride_indices_tok = ssm_state_indices.stride()
240
+
241
+ grid = (NK, NV, N * HV)
242
+ fused_recurrent_gated_delta_rule_fwd_kernel[grid](
243
+ q=q,
244
+ k=k,
245
+ v=v,
246
+ g=g,
247
+ beta=beta,
248
+ o=o,
249
+ h0=initial_state,
250
+ ht=final_state,
251
+ cu_seqlens=cu_seqlens,
252
+ ssm_state_indices=ssm_state_indices,
253
+ num_accepted_tokens=num_accepted_tokens,
254
+ scale=scale,
255
+ N=N,
256
+ T=T,
257
+ B=B,
258
+ H=H,
259
+ HV=HV,
260
+ K=K,
261
+ V=V,
262
+ BK=BK,
263
+ BV=BV,
264
+ stride_init_state_token=stride_init_state_token,
265
+ stride_final_state_token=stride_final_state_token,
266
+ stride_indices_seq=stride_indices_seq,
267
+ stride_indices_tok=stride_indices_tok,
268
+ IS_BETA_HEADWISE=beta.ndim == v.ndim,
269
+ USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,
270
+ INPLACE_FINAL_STATE=inplace_final_state,
271
+ IS_KDA=False,
272
+ SIGMOID_BETA=False,
273
+ a_log=None,
274
+ g_bias=None,
275
+ COMPUTE_GATE=False,
276
+ SAFE_GATE=True,
277
+ LOWER_BOUND=-5.0,
278
+ num_warps=num_warps,
279
+ num_stages=num_stages,
280
+ )
281
+ o = o.squeeze(0)
282
+ return o, final_state
283
+
284
+
285
+@triton.jit
286
+def fused_recurrent_gated_delta_rule_packed_decode_kernel(
287
+ mixed_qkv,
288
+ a,
289
+ b,
290
+ A_log,
291
+ dt_bias,
292
+ o,
293
+ h0,
294
+ ht,
295
+ ssm_state_indices,
296
+ scale,
297
+ stride_mixed_qkv_tok: tl.constexpr,
298
+ stride_a_tok: tl.constexpr,
299
+ stride_b_tok: tl.constexpr,
300
+ stride_init_state_token: tl.constexpr,
301
+ stride_final_state_token: tl.constexpr,
302
+ stride_indices_seq: tl.constexpr,
303
+ H: tl.constexpr,
304
+ HV: tl.constexpr,
305
+ K: tl.constexpr,
306
+ V: tl.constexpr,
307
+ BK: tl.constexpr,
308
+ BV: tl.constexpr,
309
+ SOFTPLUS_THRESHOLD: tl.constexpr,
310
+ USE_QK_L2NORM_IN_KERNEL: tl.constexpr,
311
+ SPLIT_BATCH_HEAD_GRID: tl.constexpr,
312
+):
313
+ if SPLIT_BATCH_HEAD_GRID:
314
+ i_v, i_hv, i_n = tl.program_id(0), tl.program_id(1), tl.program_id(2)
315
+ else:
316
+ i_v, i_nh = tl.program_id(0), tl.program_id(1)
317
+ i_n, i_hv = i_nh // HV, i_nh % HV
318
+ i_h = i_hv // (HV // H)
319
+
320
+ o_k = tl.arange(0, BK)
321
+ o_v = i_v * BV + tl.arange(0, BV)
322
+ mask_k = o_k < K
323
+ mask_v = o_v < V
324
+ mask_h = mask_v[:, None] & mask_k[None, :]
325
+
326
+ state_idx = tl.load(ssm_state_indices + i_n * stride_indices_seq).to(tl.int64)
327
+ p_o = o + (i_n * HV + i_hv) * V + o_v
328
+
329
+ # Skip if state index is invalid (NULL_BLOCK_ID=0)
330
+ if state_idx <= 0:
331
+ zero = tl.zeros([BV], dtype=tl.float32).to(p_o.dtype.element_ty)
332
+ tl.store(p_o, zero, mask=mask_v)
333
+ return
334
+
335
+ p_h0 = h0 + state_idx * stride_init_state_token
336
+ p_h0 = p_h0 + i_hv * V * K + o_v[:, None] * K + o_k[None, :]
337
+ b_h = tl.load(p_h0, mask=mask_h, other=0).to(tl.float32)
338
+
339
+ p_mixed = mixed_qkv + i_n * stride_mixed_qkv_tok
340
+ q_off = i_h * K + o_k
341
+ k_off = (H * K) + i_h * K + o_k
342
+ v_off = (2 * H * K) + i_hv * V + o_v
343
+ b_q = tl.load(p_mixed + q_off, mask=mask_k, other=0).to(tl.float32)
344
+ b_k = tl.load(p_mixed + k_off, mask=mask_k, other=0).to(tl.float32)
345
+ b_v = tl.load(p_mixed + v_off, mask=mask_v, other=0).to(tl.float32)
346
+
347
+ if USE_QK_L2NORM_IN_KERNEL:
348
+ b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6)
349
+ b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6)
350
+ b_q = b_q * scale
351
+
352
+ a_val = tl.load(a + i_n * stride_a_tok + i_hv).to(tl.float32)
353
+ b_val = tl.load(b + i_n * stride_b_tok + i_hv).to(tl.float32)
354
+ A_log_val = tl.load(A_log + i_hv).to(tl.float32)
355
+ dt_bias_val = tl.load(dt_bias + i_hv).to(tl.float32)
356
+ x = a_val + dt_bias_val
357
+ softplus_x = tl.where(x <= SOFTPLUS_THRESHOLD, tl.log(1.0 + tl.exp(x)), x)
358
+ g_val = -tl.exp(A_log_val) * softplus_x
359
+ beta_val = tl.sigmoid(b_val)
360
+
361
+ b_h *= exp(g_val)
362
+ b_v -= tl.sum(b_h * b_k[None, :], 1)
363
+ b_v *= beta_val
364
+ b_h += b_v[:, None] * b_k[None, :]
365
+ b_o = tl.sum(b_h * b_q[None, :], 1)
366
+ tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v)
367
+
368
+ p_ht = ht + state_idx * stride_final_state_token
369
+ p_ht = p_ht + i_hv * V * K + o_v[:, None] * K + o_k[None, :]
370
+ tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h)
371
+
372
+
373
+def fused_recurrent_gated_delta_rule_packed_decode(
374
+ mixed_qkv: torch.Tensor,
375
+ a: torch.Tensor,
376
+ b: torch.Tensor,
377
+ A_log: torch.Tensor,
378
+ dt_bias: torch.Tensor,
379
+ scale: float,
380
+ initial_state: torch.Tensor,
381
+ out: torch.Tensor,
382
+ ssm_state_indices: torch.Tensor,
383
+ use_qk_l2norm_in_kernel: bool = False,
384
+) -> tuple[torch.Tensor, torch.Tensor]:
385
+ if mixed_qkv.ndim != 2:
386
+ raise ValueError(
387
+ f"`mixed_qkv` must be a 2D tensor (got ndim={mixed_qkv.ndim})."
388
+ )
389
+ if mixed_qkv.stride(-1) != 1:
390
+ raise ValueError("`mixed_qkv` must be contiguous in the last dim.")
391
+ if a.ndim != 2 or b.ndim != 2:
392
+ raise ValueError(
393
+ f"`a` and `b` must be 2D tensors (got a.ndim={a.ndim}, b.ndim={b.ndim})."
394
+ )
395
+ if a.stride(-1) != 1 or b.stride(-1) != 1:
396
+ raise ValueError("`a`/`b` must be contiguous in the last dim.")
397
+ if A_log.ndim != 1 or dt_bias.ndim != 1:
398
+ raise ValueError("`A_log`/`dt_bias` must be 1D tensors.")
399
+ if A_log.stride(0) != 1 or dt_bias.stride(0) != 1:
400
+ raise ValueError("`A_log`/`dt_bias` must be contiguous.")
401
+ if ssm_state_indices.ndim != 1:
402
+ raise ValueError(
403
+ f"`ssm_state_indices` must be 1D for packed decode (got ndim={ssm_state_indices.ndim})."
404
+ )
405
+ if not out.is_contiguous():
406
+ raise ValueError("`out` must be contiguous.")
407
+
408
+ dev = mixed_qkv.device
409
+ if (
410
+ a.device != dev
411
+ or b.device != dev
412
+ or A_log.device != dev
413
+ or dt_bias.device != dev
414
+ or initial_state.device != dev
415
+ or out.device != dev
416
+ or ssm_state_indices.device != dev
417
+ ):
418
+ raise ValueError("All inputs must be on the same device.")
419
+
420
+ B = mixed_qkv.shape[0]
421
+ if a.shape[0] != B or b.shape[0] != B:
422
+ raise ValueError(
423
+ "Mismatched batch sizes: "
424
+ f"mixed_qkv.shape[0]={B}, a.shape[0]={a.shape[0]}, b.shape[0]={b.shape[0]}."
425
+ )
426
+ if ssm_state_indices.shape[0] != B:
427
+ raise ValueError(
428
+ f"`ssm_state_indices` must have shape [B] (got {tuple(ssm_state_indices.shape)}; expected ({B},))."
429
+ )
430
+
431
+ if initial_state.ndim != 4:
432
+ raise ValueError(
433
+ f"`initial_state` must be a 4D tensor (got ndim={initial_state.ndim})."
434
+ )
435
+ if initial_state.stride(-1) != 1:
436
+ raise ValueError("`initial_state` must be contiguous in the last dim.")
437
+ HV, V, K = initial_state.shape[-3:]
438
+ if a.shape[1] != HV or b.shape[1] != HV:
439
+ raise ValueError(
440
+ f"`a`/`b` must have shape [B, HV] with HV={HV} (got a.shape={tuple(a.shape)}, b.shape={tuple(b.shape)})."
441
+ )
442
+ if A_log.numel() != HV or dt_bias.numel() != HV:
443
+ raise ValueError(
444
+ f"`A_log` and `dt_bias` must have {HV} elements (got A_log.numel()={A_log.numel()}, dt_bias.numel()={dt_bias.numel()})."
445
+ )
446
+ if out.shape != (B, 1, HV, V):
447
+ raise ValueError(
448
+ f"`out` must have shape {(B, 1, HV, V)} (got out.shape={tuple(out.shape)})."
449
+ )
450
+
451
+ qkv_dim = mixed_qkv.shape[1]
452
+ qk_dim = qkv_dim - HV * V
453
+ if qk_dim <= 0 or qk_dim % 2 != 0:
454
+ raise ValueError(
455
+ f"Invalid packed `mixed_qkv` last dim={qkv_dim} for HV={HV}, V={V}."
456
+ )
457
+ q_dim = qk_dim // 2
458
+ if q_dim % K != 0:
459
+ raise ValueError(f"Invalid packed Q size {q_dim}: must be divisible by K={K}.")
460
+ H = q_dim // K
461
+ if H <= 0 or HV % H != 0:
462
+ raise ValueError(
463
+ f"Invalid head config inferred from mixed_qkv: H={H}, HV={HV}."
464
+ )
465
+
466
+ BK = triton.next_power_of_2(K)
467
+ if triton.cdiv(K, BK) != 1:
468
+ raise ValueError(
469
+ f"Packed decode kernel only supports NK=1 (got K={K}, BK={BK})."
470
+ )
471
+ BV = min(triton.next_power_of_2(V), 32)
472
+ num_stages = 3
473
+ num_warps = 1
474
+
475
+ stride_mixed_qkv_tok = mixed_qkv.stride(0)
476
+ stride_a_tok = a.stride(0)
477
+ stride_b_tok = b.stride(0)
478
+ stride_init_state_token = initial_state.stride(0)
479
+ stride_final_state_token = initial_state.stride(0)
480
+ stride_indices_seq = ssm_state_indices.stride(0)
481
+
482
+ NV = triton.cdiv(V, BV)
483
+ # CUDA limits grid Y/Z dimensions to 65535.
484
+ split_batch_head_grid = B * HV > 65535
485
+ grid = (NV, HV, B) if split_batch_head_grid else (NV, B * HV)
486
+ fused_recurrent_gated_delta_rule_packed_decode_kernel[grid](
487
+ mixed_qkv=mixed_qkv,
488
+ a=a,
489
+ b=b,
490
+ A_log=A_log,
491
+ dt_bias=dt_bias,
492
+ o=out,
493
+ h0=initial_state,
494
+ ht=initial_state,
495
+ ssm_state_indices=ssm_state_indices,
496
+ scale=scale,
497
+ stride_mixed_qkv_tok=stride_mixed_qkv_tok,
498
+ stride_a_tok=stride_a_tok,
499
+ stride_b_tok=stride_b_tok,
500
+ stride_init_state_token=stride_init_state_token,
501
+ stride_final_state_token=stride_final_state_token,
502
+ stride_indices_seq=stride_indices_seq,
503
+ H=H,
504
+ HV=HV,
505
+ K=K,
506
+ V=V,
507
+ BK=BK,
508
+ BV=BV,
509
+ SOFTPLUS_THRESHOLD=20.0,
510
+ USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,
511
+ SPLIT_BATCH_HEAD_GRID=split_batch_head_grid,
512
+ num_warps=num_warps,
513
+ num_stages=num_stages,
514
+ )
515
+ return out, initial_state
516
+
517
+
518
+class FusedRecurrentFunction(torch.autograd.Function):
519
+ @staticmethod
520
+ def forward(
521
+ ctx,
522
+ q: torch.Tensor,
523
+ k: torch.Tensor,
524
+ v: torch.Tensor,
525
+ g: torch.Tensor,
526
+ beta: torch.Tensor,
527
+ scale: float,
528
+ initial_state: torch.Tensor,
529
+ inplace_final_state: bool = True,
530
+ cu_seqlens: torch.Tensor | None = None,
531
+ ssm_state_indices: torch.Tensor | None = None,
532
+ num_accepted_tokens: torch.Tensor | None = None,
533
+ use_qk_l2norm_in_kernel: bool = False,
534
+ ):
535
+ o, final_state = fused_recurrent_gated_delta_rule_fwd(
536
+ q=q.contiguous(),
537
+ k=k.contiguous(),
538
+ v=v.contiguous(),
539
+ g=g.contiguous(),
540
+ beta=beta.contiguous(),
541
+ scale=scale,
542
+ initial_state=initial_state,
543
+ inplace_final_state=inplace_final_state,
544
+ cu_seqlens=cu_seqlens,
545
+ ssm_state_indices=ssm_state_indices,
546
+ num_accepted_tokens=num_accepted_tokens,
547
+ use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
548
+ )
549
+
550
+ return o, final_state
551
+
552
+
553
+def fused_recurrent_gated_delta_rule(
554
+ q: torch.Tensor,
555
+ k: torch.Tensor,
556
+ v: torch.Tensor,
557
+ g: torch.Tensor,
558
+ beta: torch.Tensor = None,
559
+ scale: float = None,
560
+ initial_state: torch.Tensor = None,
561
+ inplace_final_state: bool = True,
562
+ cu_seqlens: torch.Tensor | None = None,
563
+ ssm_state_indices: torch.Tensor | None = None,
564
+ num_accepted_tokens: torch.Tensor | None = None,
565
+ use_qk_l2norm_in_kernel: bool = False,
566
+) -> tuple[torch.Tensor, torch.Tensor]:
567
+ r"""
568
+ Args:
569
+ q (torch.Tensor):
570
+ queries of shape `[B, T, H, K]`.
571
+ k (torch.Tensor):
572
+ keys of shape `[B, T, H, K]`.
573
+ v (torch.Tensor):
574
+ values of shape `[B, T, HV, V]`.
575
+ GVA is applied if `HV > H`.
576
+ g (torch.Tensor):
577
+ g (decays) of shape `[B, T, HV]`.
578
+ beta (torch.Tensor):
579
+ betas of shape `[B, T, HV]`.
580
+ scale (Optional[int]):
581
+ Scale factor for the RetNet attention scores.
582
+ If not provided, it will default to `1 / sqrt(K)`. Default: `None`.
583
+ initial_state (Optional[torch.Tensor]):
584
+ Initial state of shape `[N, HV, V, K]` for `N` input sequences.
585
+ For equal-length input sequences, `N` equals the batch size `B`.
586
+ Default: `None`.
587
+ inplace_final_state: bool:
588
+ Whether to store the final state in-place to save memory.
589
+ Default: `True`.
590
+ cu_seqlens (torch.Tensor):
591
+ Cumulative sequence lengths of shape `[N+1]` used for variable-length training,
592
+ consistent with the FlashAttention API.
593
+ ssm_state_indices (Optional[torch.Tensor]):
594
+ Indices to map the input sequences to the initial/final states.
595
+ num_accepted_tokens (Optional[torch.Tensor]):
596
+ Number of accepted tokens for each sequence during decoding.
597
+
598
+ Returns:
599
+ o (torch.Tensor):
600
+ Outputs of shape `[B, T, HV, V]`.
601
+ final_state (torch.Tensor):
602
+ Final state of shape `[N, HV, V, K]`.
603
+
604
+ Examples::
605
+ >>> import torch
606
+ >>> import torch.nn.functional as F
607
+ >>> from einops import rearrange
608
+ >>> from fla.ops.gated_delta_rule import fused_recurrent_gated_delta_rule
609
+ # inputs with equal lengths
610
+ >>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512
611
+ >>> q = torch.randn(B, T, H, K, device='cuda')
612
+ >>> k = F.normalize(torch.randn(B, T, H, K, device='cuda'), p=2, dim=-1)
613
+ >>> v = torch.randn(B, T, HV, V, device='cuda')
614
+ >>> g = F.logsigmoid(torch.rand(B, T, HV, device='cuda'))
615
+ >>> beta = torch.rand(B, T, HV, device='cuda').sigmoid()
616
+ >>> h0 = torch.randn(B, HV, V, K, device='cuda')
617
+ >>> o, ht = fused_gated_recurrent_delta_rule(
618
+ q, k, v, g, beta,
619
+ initial_state=h0,
620
+ )
621
+ # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required
622
+ >>> q, k, v, g, beta = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, g, beta))
623
+ # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected
624
+ >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.int32)
625
+ >>> o_var, ht_var = fused_gated_recurrent_delta_rule(
626
+ q, k, v, g, beta,
627
+ initial_state=h0,
628
+ cu_seqlens=cu_seqlens
629
+ )
630
+ """
631
+ if cu_seqlens is not None and q.shape[0] != 1:
632
+ raise ValueError(
633
+ f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`."
634
+ f"Please flatten variable-length inputs before processing."
635
+ )
636
+ if scale is None:
637
+ scale = k.shape[-1] ** -0.5
638
+ else:
639
+ assert scale > 0, "scale must be positive"
640
+ if beta is None:
641
+ beta = torch.ones_like(q[..., 0])
642
+ o, final_state = FusedRecurrentFunction.apply(
643
+ q,
644
+ k,
645
+ v,
646
+ g,
647
+ beta,
648
+ scale,
649
+ initial_state,
650
+ inplace_final_state,
651
+ cu_seqlens,
652
+ ssm_state_indices,
653
+ num_accepted_tokens,
654
+ use_qk_l2norm_in_kernel,
655
+ )
656
+ return o, final_state
vllm/models/glm5next/nvidia/ops/third_party/kda/kernels.py
new
+1362
@@ -0,0 +1,1362 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang
4
+# mypy: ignore-errors
5
+#
6
+# This file contains code copied from the flash-linear-attention project.
7
+# The original source code was licensed under the MIT license and included
8
+# the following copyright notice:
9
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
10
+# ruff: noqa: E501
11
+
12
+
13
+import torch
14
+
15
+from vllm.third_party.flash_linear_attention.ops.chunk_delta_h import (
16
+ chunk_gated_delta_rule_fwd_h,
17
+)
18
+from vllm.third_party.flash_linear_attention.ops.cumsum import chunk_local_cumsum
19
+from vllm.third_party.flash_linear_attention.ops.index import prepare_chunk_indices
20
+from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd
21
+from vllm.third_party.flash_linear_attention.ops.op import exp2, log
22
+from vllm.third_party.flash_linear_attention.ops.solve_tril import solve_tril
23
+from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE, is_amd
24
+from vllm.triton_utils import tl, triton
25
+from vllm.utils.math_utils import RCP_LN2, cdiv, next_power_of_2
26
+
27
+from .fused_recurrent import fused_recurrent_gated_delta_rule_fwd_kernel
28
+
29
+BT_LIST_AUTOTUNE = [32, 64, 128]
30
+NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if is_amd else [4, 8, 16, 32]
31
+
32
+
33
+def fused_recurrent_kda_fwd(
34
+ q: torch.Tensor,
35
+ k: torch.Tensor,
36
+ v: torch.Tensor,
37
+ g: torch.Tensor,
38
+ beta: torch.Tensor,
39
+ scale: float,
40
+ initial_state: torch.Tensor,
41
+ inplace_final_state: bool = True,
42
+ cu_seqlens: torch.Tensor | None = None,
43
+ ssm_state_indices: torch.Tensor | None = None,
44
+ num_accepted_tokens: torch.Tensor | None = None,
45
+ use_qk_l2norm_in_kernel: bool = False,
46
+ out: torch.Tensor | None = None,
47
+ sigmoid_beta: bool = False,
48
+ a_log: torch.Tensor | None = None,
49
+ g_bias: torch.Tensor | None = None,
50
+ compute_gate: bool = False,
51
+ lower_bound: float | None = -5.0,
52
+) -> tuple[torch.Tensor, torch.Tensor]:
53
+ B, T, H, K, V = *k.shape, v.shape[-1]
54
+ HV = v.shape[2]
55
+ N = B if cu_seqlens is None else len(cu_seqlens) - 1
56
+ BK, BV = next_power_of_2(K), min(next_power_of_2(V), 8)
57
+ NK, NV = cdiv(K, BK), cdiv(V, BV)
58
+ assert NK == 1, "NK > 1 is not supported yet"
59
+ num_stages = 3
60
+ num_warps = 1
61
+
62
+ if compute_gate:
63
+ assert a_log is not None and g_bias is not None, (
64
+ "compute_gate requires a_log and g_bias"
65
+ )
66
+ assert lower_bound is not None, (
67
+ "compute_gate implements the bounded (safe_gate) branch only"
68
+ )
69
+ a_log = a_log.reshape(-1).contiguous()
70
+ g_bias = g_bias.reshape(-1).contiguous()
71
+
72
+ if out is None:
73
+ o = torch.empty_like(k)
74
+ else:
75
+ # Caller-provided output buffer; must be layout-compatible with the
76
+ # tensor the kernel indexes (contiguous, same shape/dtype as k).
77
+ assert out.shape == k.shape and out.dtype == k.dtype
78
+ assert out.is_contiguous()
79
+ o = out
80
+ if inplace_final_state:
81
+ final_state = initial_state
82
+ else:
83
+ final_state = q.new_empty(T, HV, V, K, dtype=initial_state.dtype)
84
+
85
+ stride_init_state_token = initial_state.stride(0)
86
+ stride_final_state_token = final_state.stride(0)
87
+
88
+ if ssm_state_indices is None:
89
+ stride_indices_seq, stride_indices_tok = 1, 1
90
+ elif ssm_state_indices.ndim == 1:
91
+ stride_indices_seq, stride_indices_tok = ssm_state_indices.stride(0), 1
92
+ else:
93
+ stride_indices_seq, stride_indices_tok = ssm_state_indices.stride()
94
+
95
+ grid = (NK, NV, N * HV)
96
+ fused_recurrent_gated_delta_rule_fwd_kernel[grid](
97
+ q=q,
98
+ k=k,
99
+ v=v,
100
+ g=g,
101
+ beta=beta,
102
+ o=o,
103
+ h0=initial_state,
104
+ ht=final_state,
105
+ cu_seqlens=cu_seqlens,
106
+ ssm_state_indices=ssm_state_indices,
107
+ num_accepted_tokens=num_accepted_tokens,
108
+ scale=scale,
109
+ N=N,
110
+ T=T,
111
+ B=B,
112
+ H=H,
113
+ HV=HV,
114
+ K=K,
115
+ V=V,
116
+ BK=BK,
117
+ BV=BV,
118
+ stride_init_state_token=stride_init_state_token,
119
+ stride_final_state_token=stride_final_state_token,
120
+ stride_indices_seq=stride_indices_seq,
121
+ stride_indices_tok=stride_indices_tok,
122
+ IS_BETA_HEADWISE=beta.ndim == v.ndim,
123
+ USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,
124
+ INPLACE_FINAL_STATE=inplace_final_state,
125
+ IS_KDA=True,
126
+ SIGMOID_BETA=sigmoid_beta,
127
+ a_log=a_log,
128
+ g_bias=g_bias,
129
+ COMPUTE_GATE=compute_gate,
130
+ SAFE_GATE=True,
131
+ LOWER_BOUND=lower_bound if lower_bound is not None else -5.0,
132
+ num_warps=num_warps,
133
+ num_stages=num_stages,
134
+ )
135
+
136
+ return o, final_state
137
+
138
+
139
+def fused_recurrent_kda(
140
+ q: torch.Tensor,
141
+ k: torch.Tensor,
142
+ v: torch.Tensor,
143
+ g: torch.Tensor,
144
+ beta: torch.Tensor = None,
145
+ scale: float = None,
146
+ initial_state: torch.Tensor = None,
147
+ inplace_final_state: bool = True,
148
+ use_qk_l2norm_in_kernel: bool = True,
149
+ cu_seqlens: torch.Tensor | None = None,
150
+ ssm_state_indices: torch.LongTensor | None = None,
151
+ num_accepted_tokens: torch.Tensor | None = None,
152
+ out: torch.Tensor | None = None,
153
+ sigmoid_beta: bool = False,
154
+ a_log: torch.Tensor | None = None,
155
+ g_bias: torch.Tensor | None = None,
156
+ compute_gate: bool = False,
157
+ lower_bound: float | None = -5.0,
158
+ **kwargs,
159
+) -> tuple[torch.Tensor, torch.Tensor]:
160
+ if cu_seqlens is not None and q.shape[0] != 1:
161
+ raise ValueError(
162
+ f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`."
163
+ f"Please flatten variable-length inputs before processing."
164
+ )
165
+ if scale is None:
166
+ scale = k.shape[-1] ** -0.5
167
+
168
+ o, final_state = fused_recurrent_kda_fwd(
169
+ q=q.contiguous(),
170
+ k=k.contiguous(),
171
+ v=v.contiguous(),
172
+ g=g.contiguous(),
173
+ beta=beta.contiguous(),
174
+ scale=scale,
175
+ initial_state=initial_state,
176
+ inplace_final_state=inplace_final_state,
177
+ cu_seqlens=cu_seqlens,
178
+ ssm_state_indices=ssm_state_indices,
179
+ num_accepted_tokens=num_accepted_tokens,
180
+ use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
181
+ out=out,
182
+ sigmoid_beta=sigmoid_beta,
183
+ a_log=a_log,
184
+ g_bias=g_bias,
185
+ compute_gate=compute_gate,
186
+ lower_bound=lower_bound,
187
+ )
188
+ return o, final_state
189
+
190
+
191
+@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
192
+@triton.autotune(
193
+ configs=[
194
+ triton.Config({"BK": BK}, num_warps=num_warps, num_stages=num_stages)
195
+ for BK in [32, 64]
196
+ for num_warps in [1, 2, 4, 8]
197
+ for num_stages in [2, 3, 4]
198
+ ],
199
+ key=["BC"],
200
+)
201
+@triton.jit(do_not_specialize=["T"])
202
+def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter(
203
+ q,
204
+ k,
205
+ g,
206
+ beta,
207
+ A,
208
+ Aqk,
209
+ scale,
210
+ cu_seqlens,
211
+ chunk_indices,
212
+ T,
213
+ H: tl.constexpr,
214
+ K: tl.constexpr,
215
+ BT: tl.constexpr,
216
+ BC: tl.constexpr,
217
+ BK: tl.constexpr,
218
+ NC: tl.constexpr,
219
+ IS_VARLEN: tl.constexpr,
220
+):
221
+ i_t, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
222
+ i_b, i_h = i_bh // H, i_bh % H
223
+ i_i, i_j = i_c // NC, i_c % NC
224
+ if IS_VARLEN:
225
+ i_n, i_t = (
226
+ tl.load(chunk_indices + i_t * 2).to(tl.int32),
227
+ tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
228
+ )
229
+ bos, eos = (
230
+ tl.load(cu_seqlens + i_n).to(tl.int32),
231
+ tl.load(cu_seqlens + i_n + 1).to(tl.int32),
232
+ )
233
+ T = eos - bos
234
+ else:
235
+ bos, eos = i_b * T, i_b * T + T
236
+
237
+ if i_t * BT + i_i * BC >= T:
238
+ return
239
+ if i_i <= i_j:
240
+ return
241
+
242
+ q += (bos * H + i_h) * K
243
+ k += (bos * H + i_h) * K
244
+ g += (bos * H + i_h) * K
245
+ A += (bos * H + i_h) * BT
246
+ Aqk += (bos * H + i_h) * BT
247
+
248
+ p_b = tl.make_block_ptr(
249
+ beta + bos * H + i_h, (T,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,)
250
+ )
251
+ b_b = tl.load(p_b, boundary_check=(0,))
252
+
253
+ b_A = tl.zeros([BC, BC], dtype=tl.float32)
254
+ b_Aqk = tl.zeros([BC, BC], dtype=tl.float32)
255
+ for i_k in range(tl.cdiv(K, BK)):
256
+ p_q = tl.make_block_ptr(
257
+ q, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)
258
+ )
259
+ p_k = tl.make_block_ptr(
260
+ k, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)
261
+ )
262
+ p_g = tl.make_block_ptr(
263
+ g, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)
264
+ )
265
+ b_kt = tl.make_block_ptr(
266
+ k, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)
267
+ )
268
+ p_gk = tl.make_block_ptr(
269
+ g, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)
270
+ )
271
+
272
+ o_k = i_k * BK + tl.arange(0, BK)
273
+ m_k = o_k < K
274
+ # [BK,]
275
+ b_gn = tl.load(g + (i_t * BT + i_i * BC) * H * K + o_k, mask=m_k, other=0)
276
+ # [BC, BK]
277
+ b_g = tl.load(p_g, boundary_check=(0, 1))
278
+ b_k = tl.load(p_k, boundary_check=(0, 1)) * exp2(b_g - b_gn[None, :])
279
+ # [BK, BC]
280
+ b_gk = tl.load(p_gk, boundary_check=(0, 1))
281
+ b_kt = tl.load(b_kt, boundary_check=(0, 1))
282
+ # [BC, BC]
283
+ b_ktg = b_kt * exp2(b_gn[:, None] - b_gk)
284
+ b_A += tl.dot(b_k, b_ktg)
285
+
286
+ b_q = tl.load(p_q, boundary_check=(0, 1))
287
+ b_qg = b_q * exp2(b_g - b_gn[None, :]) * scale
288
+ b_Aqk += tl.dot(b_qg, b_ktg)
289
+
290
+ b_A *= b_b[:, None]
291
+
292
+ p_A = tl.make_block_ptr(
293
+ A, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)
294
+ )
295
+ tl.store(p_A, b_A.to(A.dtype.element_ty), boundary_check=(0, 1))
296
+ p_Aqk = tl.make_block_ptr(
297
+ Aqk, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)
298
+ )
299
+ tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1))
300
+
301
+
302
+@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
303
+@triton.autotune(
304
+ configs=[triton.Config({}, num_warps=num_warps) for num_warps in [1, 2, 4, 8]],
305
+ key=["BK", "BT"],
306
+)
307
+@triton.jit(do_not_specialize=["T"])
308
+def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra(
309
+ q,
310
+ k,
311
+ g,
312
+ beta,
313
+ A,
314
+ Aqk,
315
+ scale,
316
+ cu_seqlens,
317
+ chunk_indices,
318
+ T,
319
+ H: tl.constexpr,
320
+ K: tl.constexpr,
321
+ BT: tl.constexpr,
322
+ BC: tl.constexpr,
323
+ BK: tl.constexpr,
324
+ IS_VARLEN: tl.constexpr,
325
+):
326
+ i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
327
+ i_b, i_h = i_bh // H, i_bh % H
328
+ if IS_VARLEN:
329
+ i_n, i_t = (
330
+ tl.load(chunk_indices + i_t * 2).to(tl.int32),
331
+ tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
332
+ )
333
+ bos, eos = (
334
+ tl.load(cu_seqlens + i_n).to(tl.int32),
335
+ tl.load(cu_seqlens + i_n + 1).to(tl.int32),
336
+ )
337
+ T = eos - bos
338
+ else:
339
+ bos, eos = i_b * T, i_b * T + T
340
+
341
+ if i_t * BT + i_i * BC >= T:
342
+ return
343
+
344
+ o_i = tl.arange(0, BC)
345
+ o_k = tl.arange(0, BK)
346
+ m_k = o_k < K
347
+ m_A = (i_t * BT + i_i * BC + o_i) < T
348
+ o_A = (bos + i_t * BT + i_i * BC + o_i) * H * BT + i_h * BT + i_i * BC
349
+
350
+ p_q = tl.make_block_ptr(
351
+ q + (bos * H + i_h) * K,
352
+ (T, K),
353
+ (H * K, 1),
354
+ (i_t * BT + i_i * BC, 0),
355
+ (BC, BK),
356
+ (1, 0),
357
+ )
358
+ p_k = tl.make_block_ptr(
359
+ k + (bos * H + i_h) * K,
360
+ (T, K),
361
+ (H * K, 1),
362
+ (i_t * BT + i_i * BC, 0),
363
+ (BC, BK),
364
+ (1, 0),
365
+ )
366
+ p_g = tl.make_block_ptr(
367
+ g + (bos * H + i_h) * K,
368
+ (T, K),
369
+ (H * K, 1),
370
+ (i_t * BT + i_i * BC, 0),
371
+ (BC, BK),
372
+ (1, 0),
373
+ )
374
+ b_q = tl.load(p_q, boundary_check=(0, 1))
375
+ b_k = tl.load(p_k, boundary_check=(0, 1))
376
+ b_g = tl.load(p_g, boundary_check=(0, 1))
377
+
378
+ p_b = beta + (bos + i_t * BT + i_i * BC + o_i) * H + i_h
379
+ b_k = b_k * tl.load(p_b, mask=m_A, other=0)[:, None]
380
+
381
+ p_kt = k + (bos + i_t * BT + i_i * BC) * H * K + i_h * K + o_k
382
+ p_gk = g + (bos + i_t * BT + i_i * BC) * H * K + i_h * K + o_k
383
+
384
+ for j in range(0, min(BC, T - i_t * BT - i_i * BC)):
385
+ b_kt = tl.load(p_kt, mask=m_k, other=0).to(tl.float32)
386
+ b_gk = tl.load(p_gk, mask=m_k, other=0).to(tl.float32)
387
+ b_ktg = b_kt[None, :] * exp2(b_g - b_gk[None, :])
388
+ b_A = tl.sum(b_k * b_ktg, 1)
389
+ b_A = tl.where(o_i > j, b_A, 0.0)
390
+ b_Aqk = tl.sum(b_q * b_ktg, 1)
391
+ b_Aqk = tl.where(o_i >= j, b_Aqk * scale, 0.0)
392
+ tl.store(A + o_A + j, b_A, mask=m_A)
393
+ tl.store(Aqk + o_A + j, b_Aqk, mask=m_A)
394
+ p_kt += H * K
395
+ p_gk += H * K
396
+
397
+
398
+def chunk_kda_scaled_dot_kkt_fwd(
399
+ q: torch.Tensor,
400
+ k: torch.Tensor,
401
+ gk: torch.Tensor | None = None,
402
+ beta: torch.Tensor | None = None,
403
+ scale: float | None = None,
404
+ cu_seqlens: torch.Tensor | None = None,
405
+ chunk_indices: torch.Tensor | None = None,
406
+ chunk_size: int = FLA_CHUNK_SIZE,
407
+ output_dtype: torch.dtype = torch.float32,
408
+) -> tuple[torch.Tensor, torch.Tensor]:
409
+ r"""
410
+ Compute beta * K * K^T.
411
+
412
+ Args:
413
+ k (torch.Tensor):
414
+ The key tensor of shape `[B, T, H, K]`.
415
+ beta (torch.Tensor):
416
+ The beta tensor of shape `[B, T, H]`.
417
+ gk (torch.Tensor):
418
+ The cumulative sum of the gate tensor of shape `[B, T, H, K]` applied to the key tensor. Default: `None`.
419
+ cu_seqlens (torch.Tensor):
420
+ The cumulative sequence lengths of the input tensor.
421
+ Default: None
422
+ chunk_size (int):
423
+ The chunk size. Default: 64.
424
+ output_dtype (torch.dtype):
425
+ The dtype of the output tensor. Default: `torch.float32`
426
+
427
+ Returns:
428
+ beta * K * K^T of shape `[B, T, H, BT]` where `BT` is the chunk size.
429
+ """
430
+ B, T, H, K = k.shape
431
+ assert K <= 256
432
+ BT = chunk_size
433
+ if chunk_indices is None and cu_seqlens is not None:
434
+ chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
435
+ NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
436
+
437
+ BC = min(16, BT)
438
+ NC = cdiv(BT, BC)
439
+ BK = max(next_power_of_2(K), 16)
440
+ A = torch.zeros(B, T, H, BT, device=k.device, dtype=output_dtype)
441
+ Aqk = torch.zeros(B, T, H, BT, device=k.device, dtype=output_dtype)
442
+ grid = (NT, NC * NC, B * H)
443
+ chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter[grid](
444
+ q=q,
445
+ k=k,
446
+ g=gk,
447
+ beta=beta,
448
+ A=A,
449
+ Aqk=Aqk,
450
+ scale=scale,
451
+ cu_seqlens=cu_seqlens,
452
+ chunk_indices=chunk_indices,
453
+ T=T,
454
+ H=H,
455
+ K=K,
456
+ BT=BT,
457
+ BC=BC,
458
+ NC=NC,
459
+ )
460
+
461
+ grid = (NT, NC, B * H)
462
+ chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra[grid](
463
+ q=q,
464
+ k=k,
465
+ g=gk,
466
+ beta=beta,
467
+ A=A,
468
+ Aqk=Aqk,
469
+ scale=scale,
470
+ cu_seqlens=cu_seqlens,
471
+ chunk_indices=chunk_indices,
472
+ T=T,
473
+ H=H,
474
+ K=K,
475
+ BT=BT,
476
+ BC=BC,
477
+ BK=BK,
478
+ )
479
+ return A, Aqk
480
+
481
+
482
+@triton.heuristics(
483
+ {
484
+ "STORE_QG": lambda args: args["qg"] is not None,
485
+ "STORE_KG": lambda args: args["kg"] is not None,
486
+ "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
487
+ }
488
+)
489
+@triton.autotune(
490
+ configs=[
491
+ triton.Config({}, num_warps=num_warps, num_stages=num_stages)
492
+ for num_warps in [2, 4, 8]
493
+ for num_stages in [2, 3, 4]
494
+ ],
495
+ key=["H", "K", "V", "BT", "BK", "BV", "IS_VARLEN"],
496
+)
497
+@triton.jit(do_not_specialize=["T"])
498
+def recompute_w_u_fwd_kernel(
499
+ q,
500
+ k,
501
+ qg,
502
+ kg,
503
+ v,
504
+ beta,
505
+ w,
506
+ u,
507
+ A,
508
+ gk,
509
+ cu_seqlens,
510
+ chunk_indices,
511
+ T,
512
+ H: tl.constexpr,
513
+ K: tl.constexpr,
514
+ V: tl.constexpr,
515
+ BT: tl.constexpr,
516
+ BK: tl.constexpr,
517
+ BV: tl.constexpr,
518
+ STORE_QG: tl.constexpr,
519
+ STORE_KG: tl.constexpr,
520
+ IS_VARLEN: tl.constexpr,
521
+ DOT_PRECISION: tl.constexpr,
522
+):
523
+ i_t, i_bh = tl.program_id(0), tl.program_id(1)
524
+ i_b, i_h = i_bh // H, i_bh % H
525
+ if IS_VARLEN:
526
+ i_n, i_t = (
527
+ tl.load(chunk_indices + i_t * 2).to(tl.int32),
528
+ tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
529
+ )
530
+ bos, eos = (
531
+ tl.load(cu_seqlens + i_n).to(tl.int32),
532
+ tl.load(cu_seqlens + i_n + 1).to(tl.int32),
533
+ )
534
+ T = eos - bos
535
+ else:
536
+ bos, eos = i_b * T, i_b * T + T
537
+ p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,))
538
+ b_b = tl.load(p_b, boundary_check=(0,))
539
+
540
+ p_A = tl.make_block_ptr(
541
+ A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)
542
+ )
543
+ b_A = tl.load(p_A, boundary_check=(0, 1))
544
+
545
+ for i_v in range(tl.cdiv(V, BV)):
546
+ p_v = tl.make_block_ptr(
547
+ v + (bos * H + i_h) * V,
548
+ (T, V),
549
+ (H * V, 1),
550
+ (i_t * BT, i_v * BV),
551
+ (BT, BV),
552
+ (1, 0),
553
+ )
554
+ p_u = tl.make_block_ptr(
555
+ u + (bos * H + i_h) * V,
556
+ (T, V),
557
+ (H * V, 1),
558
+ (i_t * BT, i_v * BV),
559
+ (BT, BV),
560
+ (1, 0),
561
+ )
562
+ b_v = tl.load(p_v, boundary_check=(0, 1))
563
+ b_vb = (b_v * b_b[:, None]).to(b_v.dtype)
564
+ b_u = tl.dot(b_A, b_vb, input_precision=DOT_PRECISION)
565
+ tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1))
566
+
567
+ for i_k in range(tl.cdiv(K, BK)):
568
+ p_w = tl.make_block_ptr(
569
+ w + (bos * H + i_h) * K,
570
+ (T, K),
571
+ (H * K, 1),
572
+ (i_t * BT, i_k * BK),
573
+ (BT, BK),
574
+ (1, 0),
575
+ )
576
+ p_k = tl.make_block_ptr(
577
+ k + (bos * H + i_h) * K,
578
+ (T, K),
579
+ (H * K, 1),
580
+ (i_t * BT, i_k * BK),
581
+ (BT, BK),
582
+ (1, 0),
583
+ )
584
+ b_k = tl.load(p_k, boundary_check=(0, 1))
585
+ b_kb = b_k * b_b[:, None]
586
+
587
+ p_gk = tl.make_block_ptr(
588
+ gk + (bos * H + i_h) * K,
589
+ (T, K),
590
+ (H * K, 1),
591
+ (i_t * BT, i_k * BK),
592
+ (BT, BK),
593
+ (1, 0),
594
+ )
595
+ b_gk = tl.load(p_gk, boundary_check=(0, 1))
596
+ b_kb *= exp2(b_gk)
597
+ if STORE_QG:
598
+ p_q = tl.make_block_ptr(
599
+ q + (bos * H + i_h) * K,
600
+ (T, K),
601
+ (H * K, 1),
602
+ (i_t * BT, i_k * BK),
603
+ (BT, BK),
604
+ (1, 0),
605
+ )
606
+ p_qg = tl.make_block_ptr(
607
+ qg + (bos * H + i_h) * K,
608
+ (T, K),
609
+ (H * K, 1),
610
+ (i_t * BT, i_k * BK),
611
+ (BT, BK),
612
+ (1, 0),
613
+ )
614
+ b_q = tl.load(p_q, boundary_check=(0, 1))
615
+ b_qg = b_q * exp2(b_gk)
616
+ tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1))
617
+ if STORE_KG:
618
+ last_idx = min(i_t * BT + BT, T) - 1
619
+
620
+ o_k = i_k * BK + tl.arange(0, BK)
621
+ m_k = o_k < K
622
+ b_gn = tl.load(
623
+ gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.0
624
+ )
625
+ b_kg = b_k * exp2(b_gn - b_gk)
626
+
627
+ p_kg = tl.make_block_ptr(
628
+ kg + (bos * H + i_h) * K,
629
+ (T, K),
630
+ (H * K, 1),
631
+ (i_t * BT, i_k * BK),
632
+ (BT, BK),
633
+ (1, 0),
634
+ )
635
+ tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1))
636
+
637
+ b_w = tl.dot(b_A, b_kb.to(b_k.dtype))
638
+ tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1))
639
+
640
+
641
+def recompute_w_u_fwd(
642
+ k: torch.Tensor,
643
+ v: torch.Tensor,
644
+ beta: torch.Tensor,
645
+ A: torch.Tensor,
646
+ q: torch.Tensor | None = None,
647
+ gk: torch.Tensor | None = None,
648
+ cu_seqlens: torch.Tensor | None = None,
649
+ chunk_indices: torch.Tensor | None = None,
650
+) -> tuple[torch.Tensor, torch.Tensor]:
651
+ B, T, H, K, V = *k.shape, v.shape[-1]
652
+ BT = A.shape[-1]
653
+ BK = 64
654
+ BV = 64
655
+
656
+ if chunk_indices is None and cu_seqlens is not None:
657
+ chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
658
+ NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
659
+
660
+ w = torch.empty_like(k)
661
+ u = torch.empty_like(v)
662
+ kg = torch.empty_like(k) if gk is not None else None
663
+ recompute_w_u_fwd_kernel[(NT, B * H)](
664
+ q=q,
665
+ k=k,
666
+ qg=None,
667
+ kg=kg,
668
+ v=v,
669
+ beta=beta,
670
+ w=w,
671
+ u=u,
672
+ A=A,
673
+ gk=gk,
674
+ cu_seqlens=cu_seqlens,
675
+ chunk_indices=chunk_indices,
676
+ T=T,
677
+ H=H,
678
+ K=K,
679
+ V=V,
680
+ BT=BT,
681
+ BK=BK,
682
+ BV=BV,
683
+ DOT_PRECISION="ieee",
684
+ )
685
+ return w, u, None, kg
686
+
687
+
688
+@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
689
+@triton.autotune(
690
+ configs=[
691
+ triton.Config({"BK": BK, "BV": BV}, num_warps=num_warps, num_stages=num_stages)
692
+ for BK in [32, 64]
693
+ for BV in [64, 128]
694
+ for num_warps in [2, 4, 8]
695
+ for num_stages in [2, 3, 4]
696
+ ],
697
+ key=["BT"],
698
+)
699
+@triton.jit(do_not_specialize=["T"])
700
+def chunk_gla_fwd_kernel_o(
701
+ q,
702
+ v,
703
+ g,
704
+ h,
705
+ o,
706
+ A,
707
+ cu_seqlens,
708
+ chunk_indices,
709
+ scale,
710
+ T,
711
+ H: tl.constexpr,
712
+ K: tl.constexpr,
713
+ V: tl.constexpr,
714
+ BT: tl.constexpr,
715
+ BK: tl.constexpr,
716
+ BV: tl.constexpr,
717
+ IS_VARLEN: tl.constexpr,
718
+):
719
+ i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
720
+ i_b, i_h = i_bh // H, i_bh % H
721
+ if IS_VARLEN:
722
+ i_tg = i_t
723
+ i_n, i_t = (
724
+ tl.load(chunk_indices + i_t * 2).to(tl.int32),
725
+ tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
726
+ )
727
+ bos, eos = (
728
+ tl.load(cu_seqlens + i_n).to(tl.int32),
729
+ tl.load(cu_seqlens + i_n + 1).to(tl.int32),
730
+ )
731
+ T = eos - bos
732
+ NT = tl.cdiv(T, BT)
733
+ else:
734
+ NT = tl.cdiv(T, BT)
735
+ i_tg = i_b * NT + i_t
736
+ bos, eos = i_b * T, i_b * T + T
737
+
738
+ m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :]
739
+
740
+ b_o = tl.zeros([BT, BV], dtype=tl.float32)
741
+ for i_k in range(tl.cdiv(K, BK)):
742
+ p_q = tl.make_block_ptr(
743
+ q + (bos * H + i_h) * K,
744
+ (T, K),
745
+ (H * K, 1),
746
+ (i_t * BT, i_k * BK),
747
+ (BT, BK),
748
+ (1, 0),
749
+ )
750
+ p_g = tl.make_block_ptr(
751
+ g + (bos * H + i_h) * K,
752
+ (T, K),
753
+ (H * K, 1),
754
+ (i_t * BT, i_k * BK),
755
+ (BT, BK),
756
+ (1, 0),
757
+ )
758
+ p_h = tl.make_block_ptr(
759
+ h + (i_tg * H + i_h) * K * V,
760
+ (V, K),
761
+ (K, 1),
762
+ (i_v * BV, i_k * BK),
763
+ (BV, BK),
764
+ (1, 0),
765
+ )
766
+
767
+ # [BT, BK]
768
+ b_q = tl.load(p_q, boundary_check=(0, 1))
769
+ b_q = (b_q * scale).to(b_q.dtype)
770
+ # [BT, BK]
771
+ b_g = tl.load(p_g, boundary_check=(0, 1))
772
+ # [BT, BK]
773
+ b_qg = (b_q * exp2(b_g)).to(b_q.dtype)
774
+ # [BV, BK]
775
+ b_h = tl.load(p_h, boundary_check=(0, 1))
776
+ # [BT, BV]
777
+ if i_k >= 0:
778
+ b_o += tl.dot(b_qg, tl.trans(b_h).to(b_qg.dtype))
779
+ p_v = tl.make_block_ptr(
780
+ v + (bos * H + i_h) * V,
781
+ (T, V),
782
+ (H * V, 1),
783
+ (i_t * BT, i_v * BV),
784
+ (BT, BV),
785
+ (1, 0),
786
+ )
787
+ p_o = tl.make_block_ptr(
788
+ o + (bos * H + i_h) * V,
789
+ (T, V),
790
+ (H * V, 1),
791
+ (i_t * BT, i_v * BV),
792
+ (BT, BV),
793
+ (1, 0),
794
+ )
795
+ p_A = tl.make_block_ptr(
796
+ A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)
797
+ )
798
+ # [BT, BV]
799
+ b_v = tl.load(p_v, boundary_check=(0, 1))
800
+ # [BT, BT]
801
+ b_A = tl.load(p_A, boundary_check=(0, 1))
802
+ b_A = tl.where(m_s, b_A, 0.0).to(b_v.dtype)
803
+ b_o += tl.dot(b_A, b_v, allow_tf32=False)
804
+ tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1))
805
+
806
+
807
+def chunk_gla_fwd_o_gk(
808
+ q: torch.Tensor,
809
+ v: torch.Tensor,
810
+ g: torch.Tensor,
811
+ A: torch.Tensor,
812
+ h: torch.Tensor,
813
+ o: torch.Tensor,
814
+ scale: float,
815
+ cu_seqlens: torch.Tensor | None = None,
816
+ chunk_indices: torch.Tensor | None = None,
817
+ chunk_size: int = FLA_CHUNK_SIZE,
818
+):
819
+ B, T, H, K, V = *q.shape, v.shape[-1]
820
+ BT = chunk_size
821
+
822
+ if chunk_indices is None and cu_seqlens is not None:
823
+ chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size)
824
+ NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
825
+
826
+ def grid(meta):
827
+ return (cdiv(V, meta["BV"]), NT, B * H)
828
+
829
+ chunk_gla_fwd_kernel_o[grid](
830
+ q=q,
831
+ v=v,
832
+ g=g,
833
+ h=h,
834
+ o=o,
835
+ A=A,
836
+ cu_seqlens=cu_seqlens,
837
+ chunk_indices=chunk_indices,
838
+ scale=scale,
839
+ T=T,
840
+ H=H,
841
+ K=K,
842
+ V=V,
843
+ BT=BT,
844
+ )
845
+ return o
846
+
847
+
848
+@triton.heuristics(
849
+ {
850
+ "HAS_BIAS": lambda args: args["g_bias"] is not None,
851
+ "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
852
+ }
853
+)
854
+@triton.autotune(
855
+ configs=[
856
+ triton.Config({"BD": BD}, num_warps=num_warps)
857
+ for BD in [32, 64]
858
+ for num_warps in [2, 4, 8]
859
+ ],
860
+ key=["H", "D", "BT", "IS_VARLEN"],
861
+)
862
+@triton.jit(do_not_specialize=["T"])
863
+def kda_gate_cumsum_fwd_kernel(
864
+ g,
865
+ A,
866
+ y,
867
+ g_bias,
868
+ cu_seqlens,
869
+ chunk_indices,
870
+ cumsum_scale,
871
+ beta,
872
+ threshold,
873
+ SAFE_GATE: tl.constexpr,
874
+ LOWER_BOUND: tl.constexpr,
875
+ T,
876
+ H: tl.constexpr,
877
+ D: tl.constexpr,
878
+ BT: tl.constexpr,
879
+ BD: tl.constexpr,
880
+ HAS_BIAS: tl.constexpr,
881
+ IS_VARLEN: tl.constexpr,
882
+):
883
+ i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
884
+ i_b, i_h = i_bh // H, i_bh % H
885
+ if IS_VARLEN:
886
+ i_n, i_t = (
887
+ tl.load(chunk_indices + i_t * 2).to(tl.int32),
888
+ tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
889
+ )
890
+ bos, eos = (
891
+ tl.load(cu_seqlens + i_n).to(tl.int32),
892
+ tl.load(cu_seqlens + i_n + 1).to(tl.int32),
893
+ )
894
+ T = eos - bos
895
+ else:
896
+ bos = i_b * T
897
+
898
+ p_g = tl.make_block_ptr(
899
+ g + (bos * H + i_h) * D,
900
+ (T, D),
901
+ (H * D, 1),
902
+ (i_t * BT, i_d * BD),
903
+ (BT, BD),
904
+ (1, 0),
905
+ )
906
+ p_y = tl.make_block_ptr(
907
+ y + (bos * H + i_h) * D,
908
+ (T, D),
909
+ (H * D, 1),
910
+ (i_t * BT, i_d * BD),
911
+ (BT, BD),
912
+ (1, 0),
913
+ )
914
+
915
+ b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32)
916
+ if HAS_BIAS:
917
+ o_d = i_d * BD + tl.arange(0, BD)
918
+ b_bias = tl.load(g_bias + i_h * D + o_d, mask=o_d < D, other=0.0).to(tl.float32)
919
+ b_g = b_g + b_bias[None, :]
920
+
921
+ b_a = tl.load(A + i_h).to(tl.float32)
922
+ b_a = tl.exp(b_a) if SAFE_GATE else -tl.exp(b_a)
923
+ if SAFE_GATE:
924
+ # y = lower_bound * sigmoid(exp(A) * (g + g_bias)), bounded to
925
+ # (lower_bound, 0) for safe-gate checkpoints.
926
+ b_gate = LOWER_BOUND / (1.0 + tl.exp(-(b_a * b_g)))
927
+ else:
928
+ b_g_scaled = b_g * beta
929
+ b_softplus = tl.where(
930
+ b_g_scaled > threshold,
931
+ b_g,
932
+ (1.0 / beta) * log(1.0 + tl.exp(b_g_scaled)),
933
+ )
934
+ b_gate = b_a * b_softplus
935
+
936
+ # Out-of-bounds rows (load returns 0, but softplus/bias can still make
937
+ # b_gate non-zero) participate in the dot product. They only contribute to
938
+ # out-of-bounds output rows, which are masked away by `boundary_check` on
939
+ # the store, so visible output matches unfused gate + chunk-local cumsum.
940
+ o_t = tl.arange(0, BT)
941
+ m_cumsum = tl.where(o_t[:, None] >= o_t[None, :], 1.0, 0.0)
942
+ b_y = tl.dot(m_cumsum, b_gate, allow_tf32=False) * cumsum_scale
943
+ tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1))
944
+
945
+
946
+def fused_kda_gate_chunk_cumsum(
947
+ raw_g: torch.Tensor,
948
+ A_log: torch.Tensor,
949
+ g_bias: torch.Tensor | None = None,
950
+ beta: float = 1.0,
951
+ threshold: float = 20.0,
952
+ cu_seqlens: torch.Tensor | None = None,
953
+ chunk_indices: torch.Tensor | None = None,
954
+ chunk_size: int = FLA_CHUNK_SIZE,
955
+ output_dtype: torch.dtype | None = torch.float,
956
+ safe_gate: bool = False,
957
+ lower_bound: float = -5.0,
958
+) -> torch.Tensor:
959
+ if cu_seqlens is not None:
960
+ assert raw_g.shape[0] == 1, (
961
+ "Only batch size 1 is supported when cu_seqlens are provided"
962
+ )
963
+ B, T, H, D = raw_g.shape
964
+ if chunk_indices is None and cu_seqlens is not None:
965
+ chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size)
966
+ NT = cdiv(T, chunk_size) if cu_seqlens is None else len(chunk_indices)
967
+
968
+ A_log = A_log.reshape(-1)
969
+ if g_bias is not None:
970
+ g_bias = g_bias.reshape(-1)
971
+ y = torch.empty_like(raw_g, dtype=output_dtype or raw_g.dtype)
972
+
973
+ def grid(meta):
974
+ return (cdiv(meta["D"], meta["BD"]), NT, B * H)
975
+
976
+ kda_gate_cumsum_fwd_kernel[grid](
977
+ g=raw_g,
978
+ A=A_log,
979
+ y=y,
980
+ g_bias=g_bias,
981
+ cu_seqlens=cu_seqlens,
982
+ chunk_indices=chunk_indices,
983
+ # RCP_LN2 folds in the natural-log -> log2 conversion so downstream
984
+ # exp2-based kernels reproduce exp(g). Keep this in sync with the
985
+ # `use_exp2=True` path in `_chunk_kda_fwd_with_cumulative_g`.
986
+ cumsum_scale=RCP_LN2,
987
+ beta=beta,
988
+ threshold=threshold,
989
+ SAFE_GATE=safe_gate,
990
+ LOWER_BOUND=lower_bound,
991
+ T=T,
992
+ H=H,
993
+ D=D,
994
+ BT=chunk_size,
995
+ )
996
+ return y
997
+
998
+
999
+def _chunk_kda_fwd_with_cumulative_g(
1000
+ q: torch.Tensor,
1001
+ k: torch.Tensor,
1002
+ v: torch.Tensor,
1003
+ g: torch.Tensor,
1004
+ beta: torch.Tensor,
1005
+ scale: float,
1006
+ initial_state: torch.Tensor,
1007
+ output_final_state: bool,
1008
+ cu_seqlens: torch.Tensor | None = None,
1009
+ chunk_indices: torch.Tensor | None = None,
1010
+ chunk_size: int = FLA_CHUNK_SIZE,
1011
+):
1012
+ # `g` must already be chunk-local cumulatively-summed AND scaled by
1013
+ # RCP_LN2 (so the downstream exp2-based kernels reproduce exp(g)).
1014
+ # Use `chunk_kda_fwd` or `chunk_kda_with_fused_gate_fwd` instead of
1015
+ # calling this helper directly unless that invariant is upheld.
1016
+ # the intra Aqk is kept in fp32
1017
+ # the computation has very marginal effect on the entire throughput
1018
+ A, Aqk = chunk_kda_scaled_dot_kkt_fwd(
1019
+ q=q,
1020
+ k=k,
1021
+ gk=g,
1022
+ beta=beta,
1023
+ scale=scale,
1024
+ cu_seqlens=cu_seqlens,
1025
+ chunk_indices=chunk_indices,
1026
+ chunk_size=chunk_size,
1027
+ output_dtype=torch.float32,
1028
+ )
1029
+ A = solve_tril(A=A, cu_seqlens=cu_seqlens, output_dtype=k.dtype)
1030
+ w, u, _, kg = recompute_w_u_fwd(
1031
+ k=k,
1032
+ v=v,
1033
+ beta=beta,
1034
+ A=A,
1035
+ gk=g,
1036
+ cu_seqlens=cu_seqlens,
1037
+ chunk_indices=chunk_indices,
1038
+ )
1039
+ del A
1040
+ h, v_new, final_state = chunk_gated_delta_rule_fwd_h(
1041
+ k=kg,
1042
+ w=w,
1043
+ u=u,
1044
+ gk=g,
1045
+ initial_state=initial_state,
1046
+ output_final_state=output_final_state,
1047
+ cu_seqlens=cu_seqlens,
1048
+ chunk_indices=chunk_indices,
1049
+ use_exp2=True,
1050
+ )
1051
+ del w, u, kg
1052
+ o = chunk_gla_fwd_o_gk(
1053
+ q=q,
1054
+ v=v_new,
1055
+ g=g,
1056
+ A=Aqk,
1057
+ h=h,
1058
+ o=v,
1059
+ scale=scale,
1060
+ cu_seqlens=cu_seqlens,
1061
+ chunk_indices=chunk_indices,
1062
+ chunk_size=chunk_size,
1063
+ )
1064
+ del Aqk, v_new, h
1065
+ return o, final_state
1066
+
1067
+
1068
+def chunk_kda_fwd(
1069
+ q: torch.Tensor,
1070
+ k: torch.Tensor,
1071
+ v: torch.Tensor,
1072
+ g: torch.Tensor,
1073
+ beta: torch.Tensor,
1074
+ scale: float,
1075
+ initial_state: torch.Tensor,
1076
+ output_final_state: bool,
1077
+ cu_seqlens: torch.Tensor | None = None,
1078
+):
1079
+ chunk_size = FLA_CHUNK_SIZE
1080
+ chunk_indices = (
1081
+ prepare_chunk_indices(cu_seqlens, chunk_size)
1082
+ if cu_seqlens is not None
1083
+ else None
1084
+ )
1085
+ g = chunk_local_cumsum(
1086
+ g,
1087
+ chunk_size=chunk_size,
1088
+ cu_seqlens=cu_seqlens,
1089
+ chunk_indices=chunk_indices,
1090
+ )
1091
+ # KDA evaluates cumulative gate decays with exp2. Convert from natural-log
1092
+ # space so exp(x) is preserved as exp2(x / ln(2)).
1093
+ g = g * RCP_LN2
1094
+ return _chunk_kda_fwd_with_cumulative_g(
1095
+ q=q,
1096
+ k=k,
1097
+ v=v,
1098
+ g=g,
1099
+ beta=beta,
1100
+ scale=scale,
1101
+ initial_state=initial_state,
1102
+ output_final_state=output_final_state,
1103
+ cu_seqlens=cu_seqlens,
1104
+ chunk_indices=chunk_indices,
1105
+ chunk_size=chunk_size,
1106
+ )
1107
+
1108
+
1109
+def chunk_kda_with_fused_gate_fwd(
1110
+ q: torch.Tensor,
1111
+ k: torch.Tensor,
1112
+ v: torch.Tensor,
1113
+ raw_g: torch.Tensor,
1114
+ beta: torch.Tensor,
1115
+ A_log: torch.Tensor,
1116
+ g_bias: torch.Tensor | None,
1117
+ scale: float,
1118
+ initial_state: torch.Tensor,
1119
+ output_final_state: bool,
1120
+ cu_seqlens: torch.Tensor | None = None,
1121
+ safe_gate: bool = False,
1122
+ lower_bound: float = -5.0,
1123
+):
1124
+ chunk_size = FLA_CHUNK_SIZE
1125
+ chunk_indices = (
1126
+ prepare_chunk_indices(cu_seqlens, chunk_size)
1127
+ if cu_seqlens is not None
1128
+ else None
1129
+ )
1130
+ g = fused_kda_gate_chunk_cumsum(
1131
+ raw_g,
1132
+ A_log=A_log,
1133
+ g_bias=g_bias,
1134
+ cu_seqlens=cu_seqlens,
1135
+ chunk_indices=chunk_indices,
1136
+ chunk_size=chunk_size,
1137
+ safe_gate=safe_gate,
1138
+ lower_bound=lower_bound,
1139
+ )
1140
+ return _chunk_kda_fwd_with_cumulative_g(
1141
+ q=q,
1142
+ k=k,
1143
+ v=v,
1144
+ g=g,
1145
+ beta=beta,
1146
+ scale=scale,
1147
+ initial_state=initial_state,
1148
+ output_final_state=output_final_state,
1149
+ cu_seqlens=cu_seqlens,
1150
+ chunk_indices=chunk_indices,
1151
+ chunk_size=chunk_size,
1152
+ )
1153
+
1154
+
1155
+def chunk_kda(
1156
+ q: torch.Tensor,
1157
+ k: torch.Tensor,
1158
+ v: torch.Tensor,
1159
+ g: torch.Tensor,
1160
+ beta: torch.Tensor,
1161
+ scale: float = None,
1162
+ initial_state: torch.Tensor = None,
1163
+ output_final_state: bool = False,
1164
+ use_qk_l2norm_in_kernel: bool = False,
1165
+ cu_seqlens: torch.Tensor | None = None,
1166
+ **kwargs,
1167
+):
1168
+ if scale is None:
1169
+ scale = k.shape[-1] ** -0.5
1170
+
1171
+ if use_qk_l2norm_in_kernel:
1172
+ q = l2norm_fwd(q.contiguous())
1173
+ k = l2norm_fwd(k.contiguous())
1174
+
1175
+ o, final_state = chunk_kda_fwd(
1176
+ q=q,
1177
+ k=k,
1178
+ v=v.contiguous(),
1179
+ g=g.contiguous(),
1180
+ beta=beta.contiguous(),
1181
+ scale=scale,
1182
+ initial_state=initial_state.contiguous(),
1183
+ output_final_state=output_final_state,
1184
+ cu_seqlens=cu_seqlens,
1185
+ )
1186
+ return o, final_state
1187
+
1188
+
1189
+def chunk_kda_with_fused_gate(
1190
+ q: torch.Tensor,
1191
+ k: torch.Tensor,
1192
+ v: torch.Tensor,
1193
+ raw_g: torch.Tensor,
1194
+ beta: torch.Tensor,
1195
+ A_log: torch.Tensor,
1196
+ g_bias: torch.Tensor | None,
1197
+ scale: float | None = None,
1198
+ initial_state: torch.Tensor | None = None,
1199
+ output_final_state: bool = False,
1200
+ use_qk_l2norm_in_kernel: bool = False,
1201
+ cu_seqlens: torch.Tensor | None = None,
1202
+ safe_gate: bool = False,
1203
+ lower_bound: float = -5.0,
1204
+ **kwargs,
1205
+):
1206
+ """Run chunk KDA from raw gate projection using fused gate+cumsum."""
1207
+ if scale is None:
1208
+ scale = k.shape[-1] ** -0.5
1209
+
1210
+ if use_qk_l2norm_in_kernel:
1211
+ q = l2norm_fwd(q.contiguous())
1212
+ k = l2norm_fwd(k.contiguous())
1213
+
1214
+ o, final_state = chunk_kda_with_fused_gate_fwd(
1215
+ q=q,
1216
+ k=k,
1217
+ v=v.contiguous(),
1218
+ raw_g=raw_g.contiguous(),
1219
+ beta=beta.contiguous(),
1220
+ A_log=A_log,
1221
+ g_bias=g_bias,
1222
+ scale=scale,
1223
+ initial_state=initial_state.contiguous() if initial_state is not None else None,
1224
+ output_final_state=output_final_state,
1225
+ cu_seqlens=cu_seqlens,
1226
+ safe_gate=safe_gate,
1227
+ lower_bound=lower_bound,
1228
+ )
1229
+ return o, final_state
1230
+
1231
+
1232
+@triton.autotune(
1233
+ configs=[
1234
+ triton.Config({"BT": bt}, num_warps=nw, num_stages=ns)
1235
+ for bt in BT_LIST_AUTOTUNE
1236
+ for nw in NUM_WARPS_AUTOTUNE
1237
+ for ns in [2, 3]
1238
+ ],
1239
+ key=["H", "D"],
1240
+)
1241
+@triton.jit
1242
+def kda_gate_fwd_kernel(
1243
+ g,
1244
+ A,
1245
+ y,
1246
+ g_bias,
1247
+ beta: tl.constexpr,
1248
+ threshold: tl.constexpr,
1249
+ SAFE_GATE: tl.constexpr,
1250
+ LOWER_BOUND: tl.constexpr,
1251
+ T,
1252
+ H,
1253
+ D: tl.constexpr,
1254
+ BT: tl.constexpr,
1255
+ BD: tl.constexpr,
1256
+ HAS_BIAS: tl.constexpr,
1257
+):
1258
+ i_t, i_h = tl.program_id(0), tl.program_id(1)
1259
+ n_t = i_t * BT
1260
+
1261
+ b_a = tl.load(A + i_h).to(tl.float32)
1262
+ b_a = tl.exp(b_a) if SAFE_GATE else -tl.exp(b_a)
1263
+
1264
+ stride_row = H * D
1265
+ stride_col = 1
1266
+
1267
+ g_ptr = tl.make_block_ptr(
1268
+ base=g + i_h * D,
1269
+ shape=(T, D),
1270
+ strides=(stride_row, stride_col),
1271
+ offsets=(n_t, 0),
1272
+ block_shape=(BT, BD),
1273
+ order=(1, 0),
1274
+ )
1275
+
1276
+ y_ptr = tl.make_block_ptr(
1277
+ base=y + i_h * D,
1278
+ shape=(T, D),
1279
+ strides=(stride_row, stride_col),
1280
+ offsets=(n_t, 0),
1281
+ block_shape=(BT, BD),
1282
+ order=(1, 0),
1283
+ )
1284
+
1285
+ b_g = tl.load(g_ptr, boundary_check=(0, 1)).to(tl.float32)
1286
+
1287
+ if HAS_BIAS:
1288
+ n_d = tl.arange(0, BD)
1289
+ bias_mask = n_d < D
1290
+ b_bias = tl.load(g_bias + i_h * D + n_d, mask=bias_mask, other=0.0).to(
1291
+ tl.float32
1292
+ )
1293
+ b_g = b_g + b_bias[None, :]
1294
+
1295
+ if SAFE_GATE:
1296
+ # y = lower_bound * sigmoid(exp(A) * (g + g_bias)), bounded to
1297
+ # (lower_bound, 0) for safe-gate checkpoints.
1298
+ b_y = LOWER_BOUND / (1.0 + tl.exp(-(b_a * b_g)))
1299
+ else:
1300
+ # softplus(x, beta) = (1/beta) * log(1 + exp(beta * x))
1301
+ # When beta * x > threshold, use linear approximation x
1302
+ # Use threshold to switch to linear when beta*x > threshold
1303
+ g_scaled = b_g * beta
1304
+ use_linear = g_scaled > threshold
1305
+ sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled)))
1306
+ b_y = b_a * sp
1307
+
1308
+ tl.store(y_ptr, b_y.to(y.dtype.element_ty), boundary_check=(0, 1))
1309
+
1310
+
1311
+def fused_kda_gate(
1312
+ g: torch.Tensor,
1313
+ A: torch.Tensor,
1314
+ head_k_dim: int,
1315
+ g_bias: torch.Tensor | None = None,
1316
+ beta: float = 1.0,
1317
+ threshold: float = 20.0,
1318
+ safe_gate: bool = False,
1319
+ lower_bound: float | None = -5.0,
1320
+) -> torch.Tensor:
1321
+ """
1322
+ Forward pass for KDA gate:
1323
+ input g: [..., H*D]
1324
+ param A: [H] or [1, 1, H, 1]
1325
+ beta: softplus beta parameter (softplus branch only)
1326
+ threshold: softplus threshold parameter (softplus branch only)
1327
+ safe_gate: when False (default) compute y = -exp(A)*softplus(g+g_bias);
1328
+ when True compute the bounded y = lower_bound*sigmoid(exp(A)*(g+g_bias))
1329
+ lower_bound: floor for the safe_gate branch (default -5.0)
1330
+ return : [..., H, D]
1331
+ """
1332
+ orig_shape = g.shape[:-1]
1333
+
1334
+ g = g.view(-1, g.shape[-1])
1335
+ T = g.shape[0]
1336
+ HD = g.shape[1]
1337
+ H = A.numel()
1338
+ assert H * head_k_dim == HD
1339
+
1340
+ y = torch.empty_like(g, dtype=torch.float32)
1341
+
1342
+ def grid(meta):
1343
+ return (cdiv(T, meta["BT"]), H)
1344
+
1345
+ kda_gate_fwd_kernel[grid](
1346
+ g,
1347
+ A,
1348
+ y,
1349
+ g_bias,
1350
+ beta,
1351
+ threshold,
1352
+ safe_gate,
1353
+ lower_bound if lower_bound is not None else -5.0,
1354
+ T,
1355
+ H,
1356
+ head_k_dim,
1357
+ BD=next_power_of_2(head_k_dim),
1358
+ HAS_BIAS=g_bias is not None,
1359
+ )
1360
+
1361
+ y = y.view(*orig_shape, H, head_k_dim)
1362
+ return y
vllm/multimodal/video.py
+112
@@ -140,6 +140,24 @@ class VideoLoader:
140
"""Sampling-algorithm-specific metadata adjustment hook."""
141
return source
142
143
+ @classmethod
144
+ def read_frames(
145
+ cls,
146
+ cap: "cv2.VideoCapture",
147
+ frame_idx: list[int],
148
+ total_frames_num: int,
149
+ *,
150
+ frame_recovery: bool = False,
151
+ ) -> tuple[npt.NDArray, list[int]]:
152
+ from vllm.multimodal.video_decoders.opencv import OpenCVVideoBackendMixin
153
+
154
+ return OpenCVVideoBackendMixin.read_frames(
155
+ cap,
156
+ frame_idx,
157
+ total_frames_num,
158
+ frame_recovery=frame_recovery,
159
+ )
160
+
161
@classmethod
162
@abstractmethod
163
def load_bytes(
@@ -664,6 +682,100 @@ class GLM46VVideoBackend(VideoBackend):
682
)
683
684
685
+@VIDEO_LOADER_REGISTRY.register(
686
+ "glm5next",
687
+ # Both spellings: the borrowed-config type string (``Glm5Next...``) and
688
+ # the dedicated transformers classes landing with the new checkpoint
689
+ # (``Glm5nextVideoProcessor``, matching ``Glm5nextImageProcessor``).
690
+ video_processor=("Glm5NextVideoProcessor", "Glm5nextVideoProcessor"),
691
+)
692
+class Glm5NextVideoBackend(VideoBackend):
693
+ """GLM-5.3-Flash fps-interval video backend.
694
+
695
+ Selects frames with the same ``glm_sample_frame_indices`` sampler the
696
+ processor falls back to, so only the sampled frames are
697
+ materialized. ``fps_interval`` semantics (default 2.0) with a
698
+ temporal-patch-scaled greedy walk, frame count capped at 2048, temporal
699
+ pairs kept even. Request overrides: ``fps`` -> fps interval,
700
+ ``max_frames`` -> frame cap, ``temporal_patch_size`` (default 2).
701
+ """
702
+
703
+ _SEEK_GAP_THRESHOLD: ClassVar[int] = 64
704
+
705
+ @classmethod
706
+ def compute_frames_index_to_sample(
707
+ cls,
708
+ source: VideoSourceMetadata,
709
+ target: VideoTargetMetadata,
710
+ **kwargs,
711
+ ) -> list[int]:
712
+ # Lazy import: the processor module sits behind the
713
+ # transformers_utils package init, which multimodal must not pull in.
714
+ from vllm.transformers_utils.processors.glm5next import (
715
+ glm_sample_frame_indices,
716
+ )
717
+
718
+ return glm_sample_frame_indices(
719
+ source.total_frames_num,
720
+ source.original_fps,
721
+ source.duration or 0,
722
+ target_fps=target.fps if target.fps > 0 else None,
723
+ max_frame_count=kwargs.get("max_frames"),
724
+ temporal_patch_size=kwargs.get("temporal_patch_size", 2),
725
+ )
726
+
727
+ @classmethod
728
+ def read_frames(
729
+ cls,
730
+ cap: "cv2.VideoCapture",
731
+ frame_idx: list[int],
732
+ total_frames_num: int,
733
+ *,
734
+ frame_recovery: bool = False,
735
+ ) -> tuple[npt.NDArray, list[int]]:
736
+ if frame_recovery:
737
+ return super().read_frames(
738
+ cap, frame_idx, total_frames_num, frame_recovery=frame_recovery
739
+ )
740
+
741
+ wanted = sorted(set(frame_idx))
742
+ width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
743
+ height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
744
+ frames = np.empty((len(wanted), height, width, 3), dtype=np.uint8)
745
+ valid_frame_indices: list[int] = []
746
+ current: int | None = None
747
+ for target in wanted:
748
+ gap = target - current if current is not None else None
749
+ if gap is not None and 2 <= gap <= cls._SEEK_GAP_THRESHOLD:
750
+ for _ in range(gap - 1):
751
+ if not cap.grab():
752
+ current = None
753
+ break
754
+ else:
755
+ current = target - 1
756
+ if current != target - 1:
757
+ cap.set(cv2.CAP_PROP_POS_FRAMES, target)
758
+ current = target - 1
759
+ ok, frame = cap.read()
760
+ if ok:
761
+ frames[len(valid_frame_indices)] = cv2.cvtColor(
762
+ frame, cv2.COLOR_BGR2RGB
763
+ )
764
+ valid_frame_indices.append(target)
765
+ current = target
766
+ else:
767
+ current = None
768
+
769
+ valid_num_frames = len(valid_frame_indices)
770
+ if valid_num_frames < len(wanted):
771
+ logger.warning(
772
+ "GLM video loading expected %d sampled frames but only loaded %d.",
773
+ len(wanted),
774
+ valid_num_frames,
775
+ )
776
+ return frames[:valid_num_frames], valid_frame_indices
777
+
778
+
779
@VIDEO_LOADER_REGISTRY.register(
780
"glmga",
781
video_processor="GlmgaVideoProcessor",
vllm/multimodal/video_decoders/opencv.py
+1
-1
@@ -42,7 +42,7 @@ def decode_opencv(
42
frame_idx = loader_cls.compute_frames_index_to_sample(
43
source=source, target=target, **sampling_kwargs
44
)
45
- frames, valid = OpenCVVideoBackendMixin.read_frames(
45
+ frames, valid = loader_cls.read_frames(
46
cap,
47
frame_idx,
48
total_frames_num=source.total_frames_num,
vllm/platforms/cuda.py
+37
-7
@@ -86,6 +86,7 @@ def _get_backend_priorities(
86
num_heads: int | None = None,
87
kv_cache_dtype: CacheDType | None = None,
88
use_non_causal: bool = False,
89
+ head_size: int | None = None,
90
) -> list[AttentionBackendEnum]:
91
"""Get backend priorities with lazy import to avoid circular dependency."""
92
from vllm.utils.torch_utils import is_quantized_kv_cache
@@ -133,13 +134,21 @@ def _get_backend_priorities(
134
AttentionBackendEnum.FLASHINFER_MLA_SPARSE_SM120,
135
]
136
else:
137
+ sparse_tail = [
138
+ AttentionBackendEnum.FLASH_ATTN_MLA_SPARSE,
139
+ AttentionBackendEnum.FLASHMLA_SPARSE,
140
+ ]
141
+ flashinfer_sparse = AttentionBackendEnum.FLASHINFER_MLA_SPARSE_SM90
142
+ if head_size == 512:
143
+ sparse_tail.insert(0, flashinfer_sparse)
144
+ else:
145
+ sparse_tail.append(flashinfer_sparse)
146
return [
147
AttentionBackendEnum.FLASH_ATTN_MLA,
148
AttentionBackendEnum.FLASHMLA,
149
AttentionBackendEnum.FLASHINFER_MLA,
150
AttentionBackendEnum.TRITON_MLA,
141
- AttentionBackendEnum.FLASH_ATTN_MLA_SPARSE,
142
- AttentionBackendEnum.FLASHMLA_SPARSE,
151
+ *sparse_tail,
152
]
153
else:
154
# SM100f defaults to FlashInfer for TRTLLM causal attention, but its non-causal
@@ -373,11 +382,12 @@ class CudaPlatformBase(Platform):
382
invalid_reasons: dict[AttentionBackendEnum, tuple[int, list[str]]] = {}
383
384
backend_priorities = _get_backend_priorities(
376
- attn_selector_config.use_mla,
377
- device_capability,
378
- num_heads,
379
- attn_selector_config.kv_cache_dtype,
380
- attn_selector_config.use_non_causal,
385
+ use_mla=attn_selector_config.use_mla,
386
+ device_capability=device_capability,
387
+ num_heads=num_heads,
388
+ kv_cache_dtype=attn_selector_config.kv_cache_dtype,
389
+ use_non_causal=attn_selector_config.use_non_causal,
390
+ head_size=attn_selector_config.head_size,
391
)
392
for priority, backend in enumerate(backend_priorities):
393
try:
@@ -400,6 +410,26 @@ class CudaPlatformBase(Platform):
410
411
return valid_backends_priorities, invalid_reasons
412
413
+ @classmethod
414
+ def _get_indexer_block_alignment(cls, vllm_config: VllmConfig) -> int | None:
415
+ index_kpool = getattr(
416
+ vllm_config.model_config.hf_text_config, "index_kpool", None
417
+ )
418
+ if not index_kpool or index_kpool <= 1:
419
+ return None
420
+ from vllm.utils.deep_gemm import PAGED_MQA_PAGE_SIZES
421
+
422
+ # kpool paged-MQA indexer: the storage block (block_size /
423
+ # index_kpool) is virtually split into pool pages, so block_size
424
+ # must be a multiple of index_kpool times a legal pool page.
425
+ page = min(PAGED_MQA_PAGE_SIZES)
426
+ if cls.is_device_capability_family(120):
427
+ # On sm120 the DeepGEMM paged-MQA kernel only accepts block_kv
428
+ # 64 for the fp8 indexer cache, so align to the largest pool
429
+ # page here to make the page split land on 64 not the min 32.
430
+ page = max(PAGED_MQA_PAGE_SIZES)
431
+ return index_kpool * page
432
+
433
@classmethod
434
def get_attn_backend_cls(
435
cls,
vllm/platforms/interface.py
+15
@@ -763,6 +763,17 @@ class Platform:
763
if cache_config.mamba_page_size_padded is not None:
764
cache_config.mamba_page_size_padded = shared_page
765
766
+ @classmethod
767
+ def _get_indexer_block_alignment(cls, vllm_config: "VllmConfig") -> int | None:
768
+ """Extra ``block_size`` multiple a sparse indexer needs, else ``None``.
769
+
770
+ The CUDA kpool paged-MQA indexer virtually splits each storage block
771
+ into pool pages, so ``block_size`` must be a multiple of
772
+ ``index_kpool * min(PAGED_MQA_PAGE_SIZES)`` — implemented in the CUDA
773
+ platform override. Other platforms impose no extra constraint.
774
+ """
775
+ return None
776
+
777
@classmethod
778
def _align_hybrid_block_size(
779
cls,
@@ -805,6 +816,7 @@ class Platform:
816
num_kv_heads=model_config.get_num_kv_heads(parallel_config),
817
head_size=model_config.get_head_size(),
818
dtype=kv_cache_dtype,
819
+ cache_dtype_str=cache_config.cache_dtype,
820
kv_quant_mode=kv_quant_mode,
821
).page_size_bytes
822
elif cache_config.cache_dtype.startswith("turboquant_"):
@@ -912,6 +924,9 @@ class Platform:
924
mamba_page_size,
925
kernel_block_alignment_size * attn_page_size_1_token,
926
)
927
+ indexer_align = cls._get_indexer_block_alignment(vllm_config)
928
+ if indexer_align:
929
+ attn_block_size = indexer_align * cdiv(attn_block_size, indexer_align)
930
931
if cache_config.block_size < attn_block_size:
932
cache_config.block_size = attn_block_size
vllm/transformers_utils/config.py
+3
@@ -104,6 +104,9 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
104
k3_dspark="K3DSparkConfig",
105
funaudiochat="FunAudioChatConfig",
106
granite4_vision="Granite4VisionConfig",
107
+ glm5_next="Glm5NextConfig",
108
+ glm5_next_text="Glm5NextTextConfig",
109
+ glm5_next_vision="Glm5NextVisionConfig",
110
hyperclovax="HyperCLOVAXConfig",
111
hy_v3="HYV3Config",
112
hy_v4="HYV4Config",
vllm/transformers_utils/configs/__init__.py
+6
@@ -40,6 +40,9 @@ _CLASS_TO_MODULE: dict[str, str] = {
40
"FunAudioChatConfig": "vllm.transformers_utils.configs.funaudiochat",
41
"FunAudioChatAudioEncoderConfig": "vllm.transformers_utils.configs.funaudiochat",
42
"Granite4VisionConfig": "vllm.transformers_utils.configs.granite4_vision",
43
+ "Glm5NextConfig": "vllm.transformers_utils.configs.glm5_next",
44
+ "Glm5NextTextConfig": "vllm.transformers_utils.configs.glm5_next",
45
+ "Glm5NextVisionConfig": "vllm.transformers_utils.configs.glm5_next",
46
"HYV3Config": "vllm.transformers_utils.configs.hy_v3",
47
"HYV4Config": "vllm.transformers_utils.configs.hy_v4",
48
"HyperCLOVAXConfig": "vllm.transformers_utils.configs.hyperclovax",
@@ -133,6 +136,9 @@ __all__ = [
136
"FunAudioChatConfig",
137
"FunAudioChatAudioEncoderConfig",
138
"Granite4VisionConfig",
139
+ "Glm5NextConfig",
140
+ "Glm5NextTextConfig",
141
+ "Glm5NextVisionConfig",
142
"HYV3Config",
143
"HYV4Config",
144
"HyperCLOVAXConfig",
vllm/transformers_utils/configs/glm5_next.py
new
+420
@@ -0,0 +1,420 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from transformers.configuration_utils import PretrainedConfig
5
+
6
+from vllm.logger import init_logger
7
+
8
+logger = init_logger(__name__)
9
+
10
+
11
+class Glm5NextTextConfig(PretrainedConfig):
12
+ model_type = "glm5_next_text"
13
+ base_config_key = "text_config"
14
+ keys_to_ignore_at_inference = ["past_key_values"]
15
+
16
+ def __init__(
17
+ self,
18
+ model_type="glm5_next_text",
19
+ vocab_size: int = 154880,
20
+ hidden_size: int = 4096,
21
+ head_dim: int | None = None,
22
+ intermediate_size: int = 12288,
23
+ num_hidden_layers: int = 45,
24
+ num_attention_heads: int = 64,
25
+ num_key_value_heads: int | None = None,
26
+ hidden_act: str = "silu",
27
+ rms_norm_eps: float = 1e-5,
28
+ pad_token_id: int | None = 151329,
29
+ bos_token_id: int | None = None,
30
+ eos_token_id: int | list[int] | None = None,
31
+ rope_parameters: dict | None = None,
32
+ max_position_embeddings: int = 1048576,
33
+ tie_word_embeddings: bool = False,
34
+ moe_intermediate_size: int = 2048,
35
+ moe_renormalize: bool = True,
36
+ scoring_func: str = "sigmoid",
37
+ n_routed_experts: int | None = 288,
38
+ num_experts_per_token: int = 7,
39
+ n_shared_experts: int = 1,
40
+ routed_scaling_factor: float = 2.5,
41
+ topk_method: str | None = None,
42
+ first_k_dense_replace: int = 0,
43
+ moe_layer_freq: int = 1,
44
+ use_grouped_topk: bool = True,
45
+ n_group: int = 1,
46
+ topk_group: int = 1,
47
+ mla: bool = True,
48
+ q_lora_rank: int | None = 1536,
49
+ kv_lora_rank: int | None = 512,
50
+ qk_nope_head_dim: int = 256,
51
+ qk_rope_head_dim: int = 0,
52
+ v_head_dim: int | None = 256,
53
+ mla_nope: bool | None = True,
54
+ num_nextn_predict_layers: int = 1,
55
+ # Per-layer layout: "linear_attention" | "deepseek_sparse_attention"
56
+ layer_types: list[str] | None = None,
57
+ # Per-layer MLP: "dense" | "sparse"
58
+ mlp_layer_types: list[str] | None = None,
59
+ # Linear-attention (KDA) head config (flattened from the old
60
+ # linear_attn_config dict).
61
+ linear_head_dim: int = 128,
62
+ linear_num_heads: int = 64,
63
+ linear_conv_kernel_dim: int = 4,
64
+ linear_lower_bound: float = -5.0,
65
+ index_head_dim: int | None = None,
66
+ index_topk: int | None = None,
67
+ index_n_heads: int | None = None,
68
+ index_dsa_use_layernorm: bool = True,
69
+ index_kpool_compress: bool = True,
70
+ # Every ``index_kpool`` indexer K entries pool into one stored entry
71
+ # (compress_ratio). The 300B checkpoint ships 4; topk runs at pool
72
+ # granularity (select_k = index_topk // index_kpool).
73
+ index_kpool: int | None = 4,
74
+ index_kpool_always_select_tail: bool = True,
75
+ indexer_rope_interleave: bool = False,
76
+ mhc: bool | None = True,
77
+ mhc_num_residual_streams: int = 4,
78
+ hc_eps: float | None = 1e-06,
79
+ mhc_tau: float = 0.05,
80
+ hres_vwnstyle: bool | None = True,
81
+ mhc_no_norm_weight: bool | None = False,
82
+ mhc_sinkhorn_iterations: int | None = 20,
83
+ mhc_post_mult_value: float | None = 2.0,
84
+ swiglu_limit: float | None = None,
85
+ logit_scale: float = 1.0,
86
+ **kwargs,
87
+ ):
88
+ # Preserve checkpoint field names and local aliases because their
89
+ # consumers use different spellings.
90
+ num_experts_per_token = kwargs.get("num_experts_per_tok", num_experts_per_token)
91
+ moe_renormalize = kwargs.get("norm_topk_prob", moe_renormalize)
92
+ mhc_num_residual_streams = kwargs.get("hc_mult", mhc_num_residual_streams)
93
+ mhc_sinkhorn_iterations = kwargs.get(
94
+ "hc_sinkhorn_iters", mhc_sinkhorn_iterations
95
+ )
96
+ # Checkpoint ships ``mla_use_nope`` (not ``mla_nope``); without this
97
+ # alias self.mla_nope silently stays at the param default.
98
+ mla_nope = kwargs.get("mla_use_nope", mla_nope)
99
+ # Checkpoints ship the KDA head config as the ``linear_attn_config``
100
+ # dict (head_dim / num_heads / short_conv_kernel_size /
101
+ # gate_lower_bound) rather than the flattened top-level fields; fold it
102
+ # in so the trained values are read instead of the param defaults
103
+ # (which only match this checkpoint by coincidence).
104
+ linear_cfg = kwargs.get("linear_attn_config") or {}
105
+ if linear_cfg:
106
+ linear_head_dim = linear_cfg.get("head_dim", linear_head_dim)
107
+ linear_num_heads = linear_cfg.get("num_heads", linear_num_heads)
108
+ linear_conv_kernel_dim = linear_cfg.get(
109
+ "short_conv_kernel_size", linear_conv_kernel_dim
110
+ )
111
+ linear_lower_bound = linear_cfg.get("gate_lower_bound", linear_lower_bound)
112
+
113
+ if index_topk is not None:
114
+ if index_dsa_use_layernorm is not True:
115
+ raise NotImplementedError(
116
+ "GLM-5.3 sparse indexer requires index_dsa_use_layernorm=True"
117
+ )
118
+ if index_kpool_compress is not True:
119
+ raise NotImplementedError(
120
+ "GLM-5.3 sparse indexer requires index_kpool_compress=True"
121
+ )
122
+ if index_kpool_always_select_tail is not True:
123
+ raise NotImplementedError(
124
+ "GLM-5.3 sparse indexer requires "
125
+ "index_kpool_always_select_tail=True"
126
+ )
127
+
128
+ if mhc:
129
+ if hres_vwnstyle is not True:
130
+ raise NotImplementedError("GLM-5.3 mHC requires hres_vwnstyle=True")
131
+ if mhc_no_norm_weight not in (False, None):
132
+ raise NotImplementedError(
133
+ "GLM-5.3 mHC requires mhc_no_norm_weight=False"
134
+ )
135
+
136
+ self.model_type = model_type
137
+ self.vocab_size = vocab_size
138
+ self.hidden_size = hidden_size
139
+ self.head_dim = (
140
+ head_dim if head_dim is not None else hidden_size // num_attention_heads
141
+ )
142
+ self.intermediate_size = intermediate_size
143
+ self.num_hidden_layers = num_hidden_layers
144
+ self.num_attention_heads = num_attention_heads
145
+
146
+ # for backward compatibility
147
+ if num_key_value_heads is None:
148
+ num_key_value_heads = num_attention_heads
149
+
150
+ self.num_key_value_heads = num_key_value_heads
151
+ self.hidden_act = hidden_act
152
+ self.rms_norm_eps = rms_norm_eps
153
+ self.max_position_embeddings = max_position_embeddings
154
+ self.rope_parameters = rope_parameters
155
+
156
+ # mla config
157
+ self.mla = mla
158
+ self.q_lora_rank = q_lora_rank
159
+ self.kv_lora_rank = kv_lora_rank
160
+ self.qk_nope_head_dim = qk_nope_head_dim
161
+ self.qk_rope_head_dim = qk_rope_head_dim
162
+ self.v_head_dim = v_head_dim
163
+ self.mla_nope = mla_nope
164
+ # moe config
165
+ self.n_routed_experts = n_routed_experts
166
+ self.num_experts_per_token = num_experts_per_token
167
+ self.moe_renormalize = moe_renormalize
168
+ self.n_shared_experts = n_shared_experts
169
+ self.routed_scaling_factor = routed_scaling_factor
170
+ self.topk_method = topk_method
171
+ self.scoring_func = scoring_func
172
+ assert self.scoring_func in ("softmax", "sigmoid")
173
+ self.moe_intermediate_size = moe_intermediate_size
174
+ self.first_k_dense_replace = first_k_dense_replace
175
+ self.moe_layer_freq = moe_layer_freq
176
+ self.use_grouped_topk = use_grouped_topk
177
+ self.n_group = n_group
178
+ self.topk_group = topk_group
179
+ self.num_nextn_predict_layers = num_nextn_predict_layers
180
+
181
+ # Per-layer attention / MLP layout. Normalize mlp_layer_types from
182
+ # first_k_dense_replace when the new-schema field is absent so layer
183
+ # construction sees a consistent layout (mirrors cohere2_moe).
184
+ self.layer_types = layer_types
185
+ if mlp_layer_types is None:
186
+ n = self.num_hidden_layers
187
+ if first_k_dense_replace is not None:
188
+ mlp_layer_types = ["dense"] * first_k_dense_replace + ["sparse"] * (
189
+ n - first_k_dense_replace
190
+ )
191
+ else:
192
+ mlp_layer_types = ["sparse"] * n
193
+ self.mlp_layer_types = mlp_layer_types
194
+
195
+ # Linear-attention (KDA) head config.
196
+ self.linear_head_dim = linear_head_dim
197
+ self.linear_num_heads = linear_num_heads
198
+ self.linear_conv_kernel_dim = linear_conv_kernel_dim
199
+ self.linear_lower_bound = linear_lower_bound
200
+
201
+ # dsa index config
202
+ self.index_head_dim = index_head_dim
203
+ self.index_topk = index_topk
204
+ self.index_n_heads = index_n_heads
205
+ self.index_dsa_use_layernorm = index_dsa_use_layernorm
206
+ self.index_kpool_compress = index_kpool_compress
207
+ self.index_kpool = index_kpool
208
+ self.index_kpool_always_select_tail = index_kpool_always_select_tail
209
+ self.indexer_rope_interleave = indexer_rope_interleave
210
+
211
+ # mhc config
212
+ self.mhc = mhc
213
+ self.mhc_num_residual_streams = mhc_num_residual_streams
214
+ self.mhc_tau = mhc_tau
215
+ self.hres_vwnstyle = hres_vwnstyle
216
+ self.hc_eps = hc_eps
217
+ self.mhc_no_norm_weight = mhc_no_norm_weight
218
+ self.mhc_sinkhorn_iterations = mhc_sinkhorn_iterations
219
+ self.mhc_post_mult_value = mhc_post_mult_value
220
+
221
+ self.swiglu_limit = swiglu_limit
222
+ self.logit_scale = logit_scale
223
+
224
+ super().__init__(
225
+ pad_token_id=pad_token_id,
226
+ bos_token_id=bos_token_id,
227
+ eos_token_id=eos_token_id,
228
+ tie_word_embeddings=tie_word_embeddings,
229
+ **kwargs,
230
+ )
231
+
232
+ @property
233
+ def is_mla(self):
234
+ return (
235
+ self.q_lora_rank is not None
236
+ or self.kv_lora_rank is not None
237
+ or self.qk_nope_head_dim is not None
238
+ or self.qk_rope_head_dim is not None
239
+ or self.v_head_dim is not None
240
+ or self.mla_nope is True
241
+ )
242
+
243
+ @property
244
+ def is_moe(self):
245
+ return self.n_routed_experts is not None
246
+
247
+ @property
248
+ def is_linear_attn(self) -> bool:
249
+ return self.layer_types is not None and any(
250
+ t == "linear_attention" for t in self.layer_types
251
+ )
252
+
253
+ def is_kda_layer(self, layer_idx: int):
254
+ return (
255
+ self.layer_types is not None
256
+ and layer_idx < len(self.layer_types)
257
+ and self.layer_types[layer_idx] == "linear_attention"
258
+ )
259
+
260
+ @property
261
+ def layers_block_type(self):
262
+ # Map the schema's per-layer types onto the block strings vLLM's hybrid
263
+ # accounting (get_num_layers_by_block_type) recognizes: linear-attention
264
+ # layers stay "linear_attention"; every other attention variant collapses
265
+ # to "attention".
266
+ if self.layer_types is None:
267
+ return ["attention"] * self.num_hidden_layers
268
+ return [
269
+ "linear_attention" if t == "linear_attention" else "attention"
270
+ for t in self.layer_types
271
+ ]
272
+
273
+
274
+class Glm5NextVisionConfig(PretrainedConfig):
275
+ model_type = "glm5_next_vision"
276
+ base_config_key = "vision_config"
277
+
278
+ def __init__(
279
+ self,
280
+ depth: int = 24,
281
+ hidden_size: int = 1024,
282
+ hidden_act: str = "silu",
283
+ image_size: int = 448,
284
+ intermediate_size: int = 4096,
285
+ num_heads: int = 16,
286
+ out_hidden_size: int = 4096,
287
+ projection_intermediate_size: int = 10240,
288
+ in_channels: int = 3,
289
+ initializer_range: float = 0.02,
290
+ patch_size: int = 14,
291
+ rms_norm_eps: float = 1e-5,
292
+ spatial_merge_size: int = 2,
293
+ temporal_patch_size: int = 2,
294
+ attention_dropout: float = 0.0,
295
+ attention_bias: bool = True,
296
+ swiglu_limit: float | None = None,
297
+ **kwargs,
298
+ ):
299
+ super().__init__(**kwargs)
300
+
301
+ self.depth = depth
302
+ self.hidden_size = hidden_size
303
+ self.hidden_act = hidden_act
304
+ self.image_size = image_size
305
+ self.intermediate_size = intermediate_size
306
+ self.num_heads = num_heads
307
+ self.out_hidden_size = out_hidden_size
308
+ # GLM-5.3-Flash merger bottleneck width (absent from the generic
309
+ # GLM-OCR vision config); the tower uses it as the PatchMerger
310
+ # context_dim instead of text_config.intermediate_size.
311
+ self.projection_intermediate_size = projection_intermediate_size
312
+ self.in_channels = in_channels
313
+ self.initializer_range = initializer_range
314
+ self.patch_size = patch_size
315
+ # GLM-5.3-Flash checkpoints ship vision_config.rms_norm_eps = 1e-5,
316
+ # but the vision tower was trained with 1e-6. Serving with 1e-5 drifts
317
+ # the RMSNorm and produces repetitive/degraded image descriptions, so
318
+ # force the trained value regardless of the checkpoint field.
319
+ self.rms_norm_eps = 1e-6
320
+ self.spatial_merge_size = spatial_merge_size
321
+ self.temporal_patch_size = temporal_patch_size
322
+ self.attention_dropout = attention_dropout
323
+ self.attention_bias = attention_bias
324
+ self.swiglu_limit = swiglu_limit
325
+
326
+
327
+class Glm5NextConfig(PretrainedConfig):
328
+ model_type = "glm5_next"
329
+ sub_configs = {
330
+ "vision_config": Glm5NextVisionConfig,
331
+ "text_config": Glm5NextTextConfig,
332
+ }
333
+ keys_to_ignore_at_inference = ["past_key_values"]
334
+
335
+ def __init__(
336
+ self,
337
+ text_config=None,
338
+ vision_config=None,
339
+ image_token_id: int = 154854,
340
+ video_token_id: int = 154855,
341
+ image_start_token_id: int = 154830,
342
+ image_end_token_id: int = 154831,
343
+ video_start_token_id: int = 154832,
344
+ video_end_token_id: int = 154833,
345
+ **kwargs,
346
+ ):
347
+ # Init super() first so base-class defaults don't clobber text-config
348
+ # values set below (PretrainedConfig has many text-related defaults
349
+ # that differ from Glm5NextTextConfig).
350
+ super().__init__(**kwargs)
351
+
352
+ if isinstance(vision_config, dict):
353
+ self.vision_config = self.sub_configs["vision_config"](**vision_config)
354
+ elif vision_config is None:
355
+ self.vision_config = self.sub_configs["vision_config"]()
356
+ else:
357
+ self.vision_config = vision_config
358
+
359
+ if isinstance(text_config, dict):
360
+ self.text_config = self.sub_configs["text_config"](**text_config)
361
+ elif text_config is None:
362
+ # Backward compatibility: a flat top-level checkpoint (no nested
363
+ # text_config) folds its text fields into Glm5NextTextConfig.
364
+ self.text_config = self.sub_configs["text_config"](**kwargs)
365
+ else:
366
+ self.text_config = text_config
367
+
368
+ self.image_token_id = image_token_id
369
+ self.video_token_id = video_token_id
370
+ self.image_start_token_id = image_start_token_id
371
+ self.image_end_token_id = image_end_token_id
372
+ self.video_start_token_id = video_start_token_id
373
+ self.video_end_token_id = video_end_token_id
374
+
375
+ # Mirror attention implementation recursively onto sub-configs.
376
+ self._attn_implementation = kwargs.pop("attn_implementation", None)
377
+
378
+ # Config-metadata fields that belong to the top-level (multimodal) config
379
+ # and must NOT be mirrored onto text_config: ``architectures`` /
380
+ # ``torch_dtype`` differ between the top-level config and the text
381
+ # sub-config, and mirroring them makes the top-level ``architectures``
382
+ # silently read back as None (PretrainedConfig initializes both to None),
383
+ # which then fails model-class resolution ("No model architectures are
384
+ # specified").
385
+ _UNMIRRORED_KEYS = [
386
+ "_name_or_path",
387
+ "model_type",
388
+ "dtype",
389
+ "torch_dtype",
390
+ "architectures",
391
+ "_attn_implementation_internal",
392
+ ]
393
+
394
+ def __setattr__(self, key, value):
395
+ unmirrored = type(self)._UNMIRRORED_KEYS
396
+ if (
397
+ (text_config := super().__getattribute__("__dict__").get("text_config"))
398
+ is not None
399
+ and key not in unmirrored
400
+ and key in text_config.__dict__
401
+ ):
402
+ setattr(text_config, key, value)
403
+ else:
404
+ super().__setattr__(key, value)
405
+
406
+ def __getattribute__(self, key):
407
+ unmirrored = type(self)._UNMIRRORED_KEYS
408
+ if (
409
+ "text_config" in super().__getattribute__("__dict__")
410
+ and key not in unmirrored
411
+ ):
412
+ text_config = super().__getattribute__("text_config")
413
+ # Forward both instance attributes AND class-defined properties/
414
+ # methods of the text config, so a flat text-only checkpoint
415
+ # (model_type "glm5_next", no nested text_config) sees is_moe /
416
+ # is_kda_layer / layers_block_type like a Glm5NextTextConfig.
417
+ if key in text_config.__dict__ or key in type(text_config).__dict__:
418
+ return getattr(text_config, key)
419
+
420
+ return super().__getattribute__(key)
vllm/transformers_utils/model_arch_config_convertor.py
+2
@@ -320,6 +320,8 @@ class ModelArchConfigConvertorBase:
320
"dots3_note",
321
"deepseek_mtp",
322
"k3_dspark",
323
+ "glm5_next",
324
+ "glm5_next_text",
325
"glm_moe_dsa",
326
"glm4_moe_lite",
327
"glm4_moe_lite_mtp",
vllm/transformers_utils/processors/__init__.py
+2
@@ -18,6 +18,7 @@ __all__ = [
18
"FireRedASR2Processor",
19
"FunASRProcessor",
20
"GLM4VProcessor",
21
+ "Glm5NextProcessor",
22
"Granite4VisionProcessor",
23
"H2OVLProcessor",
24
"Moondream3Processor",
@@ -56,6 +57,7 @@ _CLASS_TO_MODULE: dict[str, str] = {
57
"FireRedASR2Processor": "vllm.transformers_utils.processors.fireredasr2",
58
"FunASRProcessor": "vllm.transformers_utils.processors.funasr",
59
"GLM4VProcessor": "vllm.transformers_utils.processors.glm4v",
60
+ "Glm5NextProcessor": "vllm.transformers_utils.processors.glm5next",
61
"Granite4VisionProcessor": "vllm.transformers_utils.processors.granite4_vision",
62
"H2OVLProcessor": "vllm.transformers_utils.processors.h2ovl",
63
"InternVLProcessor": "vllm.transformers_utils.processors.internvl",
vllm/transformers_utils/processors/glm5next.py
new
+953
@@ -0,0 +1,953 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""vLLM-native multimodal processor for GLM-5.3-Flash."""
4
+
5
+import math
6
+
7
+import numpy as np
8
+import torch
9
+from torchvision.transforms.v2 import functional as tvF
10
+from transformers.image_processing_utils import BatchFeature
11
+from transformers.image_processing_utils_fast import (
12
+ BaseImageProcessorFast,
13
+ group_images_by_shape,
14
+ reorder_images,
15
+)
16
+from transformers.image_utils import (
17
+ OPENAI_CLIP_MEAN,
18
+ OPENAI_CLIP_STD,
19
+ ChannelDimension,
20
+ ImageInput,
21
+ PILImageResampling,
22
+ SizeDict,
23
+ get_image_size,
24
+)
25
+from transformers.models.auto.image_processing_auto import get_image_processor_config
26
+from transformers.processing_utils import (
27
+ ImagesKwargs,
28
+ MultiModalData,
29
+ ProcessingKwargs,
30
+ ProcessorMixin,
31
+ Unpack,
32
+ VideosKwargs,
33
+)
34
+from transformers.tokenization_utils_base import PreTokenizedInput, TextInput
35
+from transformers.utils import TensorType, logging
36
+from transformers.video_processing_utils import BaseVideoProcessor
37
+from transformers.video_utils import (
38
+ VideoInput,
39
+ VideoMetadata,
40
+ group_videos_by_shape,
41
+ reorder_videos,
42
+)
43
+
44
+from vllm.transformers_utils.repo_utils import get_hf_file_to_dict
45
+
46
+logger = logging.get_logger(__name__)
47
+
48
+# Cap video inputs at 30,000 vision tokens to keep encoder profiling from
49
+# starving the KV cache. Image inputs retain their checkpoint-defined budget.
50
+_MAX_VIDEO_TOKENS = 30000
51
+
52
+# Frame-sampler fallbacks (mirror the checkpoint's fps_interval=2 /
53
+# max_frame_count_dynamic=2048); used when neither the request nor the
54
+# processor config overrides them.
55
+GLM_VIDEO_DEFAULT_FPS = 2.0
56
+GLM_VIDEO_DEFAULT_MAX_FRAMES = 2048
57
+
58
+
59
+def glm_sample_frame_indices(
60
+ total_frames: int,
61
+ fps: float,
62
+ duration: float,
63
+ *,
64
+ target_fps: float | None = None,
65
+ max_frame_count: int | None = None,
66
+ temporal_patch_size: int = 2,
67
+) -> list[int]:
68
+ """GLM video frame sampling (training-reference parity).
69
+
70
+ ``target_fps`` is the ``fps_interval`` request knob. The greedy walk
71
+ advances at ``1 / (temporal_patch_size * target_fps)`` seconds, so on
72
+ frame-dense sources it collects more candidates than ``extract_t`` and
73
+ the ``> extract_t`` fixup re-spreads the picks uniformly with
74
+ ``np.linspace`` -- that fallback is the intended reference behavior, not
75
+ an accident. Short clips (fewer frames than ``extract_t``) are spread at
76
+ evenly spaced timestamps (``floor`` sampling; the linspace variant
77
+ samples frames unevenly and cost 4 points on video grounding evals).
78
+ Request overrides: ``target_fps`` -> fps interval, ``max_frame_count``
79
+ -> frame cap.
80
+ """
81
+ max_frame_idx = total_frames - 1
82
+ if not duration:
83
+ duration = (round(max_frame_idx / fps) + 1) if fps else 0
84
+ if max_frame_count is None:
85
+ max_frame_count = GLM_VIDEO_DEFAULT_MAX_FRAMES
86
+ if target_fps is None:
87
+ target_fps = GLM_VIDEO_DEFAULT_FPS
88
+
89
+ extract_t = int(duration * target_fps)
90
+ extract_t = min(extract_t, int(max_frame_count))
91
+
92
+ duration_per_frame = 1 / fps
93
+ timestamps = [i * duration_per_frame for i in range(total_frames)]
94
+ max_second = int(duration)
95
+
96
+ if total_frames < extract_t:
97
+ frame_indices = [
98
+ math.floor(_i * total_frames / extract_t) for _i in range(extract_t)
99
+ ]
100
+ else:
101
+ frame_indices = []
102
+ current_second = 0.0
103
+ inv_fps = 1 / (temporal_patch_size * target_fps)
104
+ for frame_index in range(total_frames):
105
+ if timestamps[frame_index] >= current_second:
106
+ current_second += inv_fps
107
+ frame_indices.append(frame_index)
108
+ if current_second >= max_second:
109
+ break
110
+
111
+ if len(frame_indices) < extract_t:
112
+ if len(frame_indices) == 0:
113
+ start, end = 0, max(total_frames - 1, 0)
114
+ else:
115
+ start, end = frame_indices[0], frame_indices[-1]
116
+ frame_indices = np.linspace(start, end, extract_t, dtype=int).tolist()
117
+ elif len(frame_indices) > extract_t:
118
+ frame_indices = np.linspace(0, total_frames - 1, extract_t, dtype=int).tolist()
119
+
120
+ seen, uniq = set(), []
121
+ for idx in frame_indices:
122
+ if idx not in seen:
123
+ seen.add(idx)
124
+ uniq.append(int(idx))
125
+
126
+ if len(uniq) & 1:
127
+ uniq.append(uniq[-1])
128
+
129
+ return uniq
130
+
131
+
132
+def _ceil_to_factor(value: int, factor: int) -> int:
133
+ """Round a positive integer upward to the nearest multiple of factor."""
134
+ return math.ceil(value / factor) * factor
135
+
136
+
137
+def _fit_aligned_size_within_budget(
138
+ t: int,
139
+ h: int,
140
+ w: int,
141
+ h_factor: int,
142
+ w_factor: int,
143
+ max_pixels: int,
144
+) -> tuple[int, int]:
145
+ """Largest proportional size whose upward-aligned canvas fits the budget.
146
+
147
+ Binary search on the unaligned content height; each candidate is rounded
148
+ upward to h_factor/w_factor, so the returned canvas always satisfies
149
+ ``t * aligned_h * aligned_w <= max_pixels``.
150
+ """
151
+ minimum_pixels = t * h_factor * w_factor
152
+ if max_pixels < minimum_pixels:
153
+ raise ValueError(
154
+ f"max_pixels={max_pixels} is too small. At least "
155
+ f"{minimum_pixels} pixels are required for one aligned patch."
156
+ )
157
+
158
+ low, high = 1, h
159
+ best_h, best_w = h_factor, w_factor
160
+ while low <= high:
161
+ content_h = (low + high) // 2
162
+ content_w = max(1, math.floor(w * content_h / h))
163
+ aligned_h = _ceil_to_factor(content_h, h_factor)
164
+ aligned_w = _ceil_to_factor(content_w, w_factor)
165
+ if t * aligned_h * aligned_w <= max_pixels:
166
+ best_h, best_w = aligned_h, aligned_w
167
+ low = content_h + 1
168
+ else:
169
+ high = content_h - 1
170
+ return best_h, best_w
171
+
172
+
173
+def smart_resize(
174
+ t: int,
175
+ h: int,
176
+ w: int,
177
+ t_factor: int = 1,
178
+ h_factor: int = 28,
179
+ w_factor: int = 28,
180
+ min_pixels: int = 56 * 56,
181
+ max_pixels: int = 14 * 14 * 4 * 1280,
182
+) -> tuple[int, int]:
183
+ """GLM-5.3-Flash ``smart_resize``: upward-aligned canvas under a
184
+ ``t_bar * h_bar * w_bar`` pixel budget.
185
+
186
+ Height/width always round UP to their factors (content is then padded,
187
+ never cropped or distorted); an over-budget canvas is refit by binary
188
+ search instead of one-shot square-root scaling. ``h_factor`` /
189
+ ``w_factor`` carry ``patch_expand_factor`` on top of
190
+ ``patch_size * merge_size``; ``t_factor`` is ``temporal_patch_size``. For
191
+ a still image ``t = t_factor = temporal_patch_size`` so ``t_bar =
192
+ temporal_patch_size``.
193
+ """
194
+ if min(t, h, w, t_factor, h_factor, w_factor) <= 0:
195
+ raise ValueError("Image dimensions and alignment factors must be positive.")
196
+ if min_pixels <= 0 or max_pixels <= 0:
197
+ raise ValueError("min_pixels and max_pixels must be positive.")
198
+ if min_pixels > max_pixels:
199
+ raise ValueError("min_pixels must be less than or equal to max_pixels.")
200
+
201
+ t_bar = max(t_factor, round(t / t_factor) * t_factor)
202
+ h_bar = _ceil_to_factor(h, h_factor)
203
+ w_bar = _ceil_to_factor(w, w_factor)
204
+
205
+ if t_bar * h_bar * w_bar > max_pixels:
206
+ h_bar, w_bar = _fit_aligned_size_within_budget(
207
+ t=t_bar,
208
+ h=h,
209
+ w=w,
210
+ h_factor=h_factor,
211
+ w_factor=w_factor,
212
+ max_pixels=max_pixels,
213
+ )
214
+ elif t_bar * h_bar * w_bar < min_pixels:
215
+ beta = math.sqrt(min_pixels / (t * h * w))
216
+ h_bar = _ceil_to_factor(max(1, math.ceil(h * beta)), h_factor)
217
+ w_bar = _ceil_to_factor(max(1, math.ceil(w * beta)), w_factor)
218
+
219
+ # Alignment can push a candidate slightly over a tight max_pixels
220
+ # budget. Refit it when that happens.
221
+ if t_bar * h_bar * w_bar > max_pixels:
222
+ h_bar, w_bar = _fit_aligned_size_within_budget(
223
+ t=t_bar,
224
+ h=h,
225
+ w=w,
226
+ h_factor=h_factor,
227
+ w_factor=w_factor,
228
+ max_pixels=max_pixels,
229
+ )
230
+
231
+ return h_bar, w_bar
232
+
233
+
234
+def _get_pad_content_size(
235
+ image_height: int,
236
+ image_width: int,
237
+ canvas_height: int,
238
+ canvas_width: int,
239
+ allow_upscale: bool = False,
240
+) -> tuple[int, int]:
241
+ """Aspect-ratio-preserving content size that fits the canvas.
242
+
243
+ Oversized images are shrunk proportionally. Small images are enlarged
244
+ only when ``allow_upscale``. Padding is applied after the resize.
245
+ """
246
+ scale = min(canvas_height / image_height, canvas_width / image_width)
247
+ if not allow_upscale:
248
+ scale = min(1.0, scale)
249
+ content_height = max(1, min(canvas_height, math.floor(image_height * scale)))
250
+ content_width = max(1, min(canvas_width, math.floor(image_width * scale)))
251
+ return content_height, content_width
252
+
253
+
254
+def _resize_or_pad(
255
+ stacked_images: torch.Tensor,
256
+ target_height: int,
257
+ target_width: int,
258
+ resize_mode: str,
259
+ resample: "PILImageResampling | tvF.InterpolationMode | int | None",
260
+ resize,
261
+ allow_upscale: bool = False,
262
+) -> torch.Tensor:
263
+ """Resize onto the aligned canvas, or keep the aspect ratio and
264
+ zero-pad the right/bottom sides (``resize_mode="pad"``)."""
265
+ height, width = stacked_images.shape[-2:]
266
+
267
+ if resize_mode == "resize":
268
+ return resize(
269
+ stacked_images,
270
+ size=SizeDict(height=target_height, width=target_width),
271
+ resample=resample,
272
+ )
273
+
274
+ if resize_mode != "pad":
275
+ raise ValueError("resize_mode must be either 'resize' or 'pad'.")
276
+
277
+ content_height, content_width = _get_pad_content_size(
278
+ image_height=height,
279
+ image_width=width,
280
+ canvas_height=target_height,
281
+ canvas_width=target_width,
282
+ allow_upscale=allow_upscale,
283
+ )
284
+
285
+ if (content_height, content_width) != (height, width):
286
+ stacked_images = resize(
287
+ stacked_images,
288
+ size=SizeDict(height=content_height, width=content_width),
289
+ resample=resample,
290
+ )
291
+
292
+ # torchvision padding order: [left, top, right, bottom] -> pad only the
293
+ # right and bottom sides.
294
+ return tvF.pad(
295
+ stacked_images,
296
+ padding=[0, 0, target_width - content_width, target_height - content_height],
297
+ fill=0,
298
+ )
299
+
300
+
301
+def _pixel_budget(
302
+ min_image_tokens: int | None,
303
+ max_image_tokens: int | None,
304
+ patch_size: int,
305
+ merge_size: int,
306
+ temporal_patch_size: int,
307
+) -> tuple[int, int]:
308
+ """(min_pixels, max_pixels) from the token bounds of
309
+ ``processor_config.json``; one vision token covers
310
+ ``temporal_patch_size * (patch_size * merge_size) ** 2`` pixels."""
311
+ if min_image_tokens is None or max_image_tokens is None:
312
+ raise ValueError(
313
+ "min_image_tokens and max_image_tokens must be provided by "
314
+ "processor_config.json (or per-call kwargs)."
315
+ )
316
+ factor = temporal_patch_size * (patch_size * merge_size) ** 2
317
+ return min_image_tokens * factor, max_image_tokens * factor
318
+
319
+
320
+class Glm5NextImageProcessorKwargs(ImagesKwargs, total=False): # type: ignore[call-arg]
321
+ patch_size: int | None
322
+ temporal_patch_size: int | None
323
+ merge_size: int | None
324
+ patch_expand_factor: int | None
325
+ resize_mode: str | None
326
+ min_image_tokens: int | None
327
+ max_image_tokens: int | None
328
+
329
+
330
+class Glm5NextImageProcessor(BaseImageProcessorFast):
331
+ """Fast torchvision image processor for GLM-5.3-Flash.
332
+
333
+ ``patch_expand_factor`` multiplies into the ``smart_resize`` spatial
334
+ factor ``patch_size * merge_size``. ``resize_mode`` picks the geometry:
335
+ ``"pad"`` (default) preserves the aspect ratio and zero-pads the
336
+ right/bottom of the upward-aligned canvas, ``"resize"`` stretches onto
337
+ it. Defaults mirror the checkpoint's ``image_processor`` config.
338
+ """
339
+
340
+ do_resize = True
341
+ resample = PILImageResampling.BICUBIC
342
+ size = {"longest_edge": 1} # unused: budgets come from the token bounds
343
+ do_rescale = True
344
+ do_normalize = True
345
+ image_mean = OPENAI_CLIP_MEAN
346
+ image_std = OPENAI_CLIP_STD
347
+ do_convert_rgb = True
348
+ patch_size = 14
349
+ temporal_patch_size = 2
350
+ merge_size = 2
351
+ patch_expand_factor = 1
352
+ resize_mode = "pad"
353
+ min_image_tokens = 16
354
+ max_image_tokens = 8000
355
+ valid_kwargs = Glm5NextImageProcessorKwargs
356
+ model_input_names = ["pixel_values", "image_grid_thw"]
357
+
358
+ def _preprocess(
359
+ self,
360
+ images: list[torch.Tensor],
361
+ do_resize: bool,
362
+ size: SizeDict,
363
+ resample: "PILImageResampling | tvF.InterpolationMode | int | None",
364
+ do_rescale: bool,
365
+ rescale_factor: float,
366
+ do_normalize: bool,
367
+ image_mean: float | list[float] | None,
368
+ image_std: float | list[float] | None,
369
+ patch_size: int,
370
+ temporal_patch_size: int,
371
+ merge_size: int,
372
+ patch_expand_factor: int,
373
+ resize_mode: str | None,
374
+ min_image_tokens: int | None,
375
+ max_image_tokens: int | None,
376
+ disable_grouping: bool | None,
377
+ return_tensors: str | TensorType | None,
378
+ **kwargs,
379
+ ) -> BatchFeature:
380
+ resize_mode = resize_mode if resize_mode is not None else self.resize_mode
381
+ min_pixels, max_pixels = _pixel_budget(
382
+ min_image_tokens if min_image_tokens is not None else self.min_image_tokens,
383
+ max_image_tokens if max_image_tokens is not None else self.max_image_tokens,
384
+ patch_size,
385
+ merge_size,
386
+ temporal_patch_size,
387
+ )
388
+ grouped_images, grouped_images_index = group_images_by_shape(
389
+ images, disable_grouping=disable_grouping
390
+ )
391
+ resized_images_grouped = {}
392
+ for shape, stacked_images in grouped_images.items():
393
+ height, width = stacked_images.shape[-2:]
394
+ if do_resize:
395
+ resized_height, resized_width = smart_resize(
396
+ t=temporal_patch_size,
397
+ h=height,
398
+ w=width,
399
+ t_factor=temporal_patch_size,
400
+ h_factor=patch_size * merge_size * patch_expand_factor,
401
+ w_factor=patch_size * merge_size * patch_expand_factor,
402
+ min_pixels=min_pixels,
403
+ max_pixels=max_pixels,
404
+ )
405
+ stacked_images = _resize_or_pad(
406
+ stacked_images,
407
+ target_height=resized_height,
408
+ target_width=resized_width,
409
+ resize_mode=resize_mode,
410
+ resample=resample,
411
+ resize=self.resize,
412
+ allow_upscale=(temporal_patch_size * height * width < min_pixels),
413
+ )
414
+ resized_images_grouped[shape] = stacked_images
415
+
416
+ resized_images = reorder_images(resized_images_grouped, grouped_images_index)
417
+
418
+ grouped_images, grouped_images_index = group_images_by_shape(
419
+ resized_images, disable_grouping=disable_grouping
420
+ )
421
+ processed_images_grouped = {}
422
+ processed_grids = {}
423
+
424
+ for shape, stacked_images in grouped_images.items():
425
+ resized_height, resized_width = stacked_images.shape[-2:]
426
+
427
+ patches = self.rescale_and_normalize(
428
+ stacked_images,
429
+ do_rescale,
430
+ rescale_factor,
431
+ do_normalize,
432
+ image_mean,
433
+ image_std,
434
+ )
435
+ if patches.ndim == 4: # (B, C, H, W)
436
+ patches = patches.unsqueeze(1) # (B, T=1, C, H, W)
437
+
438
+ if patches.shape[1] % temporal_patch_size != 0:
439
+ repeats = patches[:, -1:].repeat(
440
+ 1,
441
+ temporal_patch_size - (patches.shape[1] % temporal_patch_size),
442
+ 1,
443
+ 1,
444
+ 1,
445
+ )
446
+ patches = torch.cat([patches, repeats], dim=1)
447
+
448
+ batch_size, t_len, channel = patches.shape[:3]
449
+ grid_t = t_len // temporal_patch_size
450
+ grid_h, grid_w = resized_height // patch_size, resized_width // patch_size
451
+
452
+ patches = patches.view(
453
+ batch_size,
454
+ grid_t,
455
+ temporal_patch_size,
456
+ channel,
457
+ grid_h // merge_size,
458
+ merge_size,
459
+ patch_size,
460
+ grid_w // merge_size,
461
+ merge_size,
462
+ patch_size,
463
+ )
464
+ # (B, grid_t, gh, gw, mh, mw, C, tp, ph, pw)
465
+ patches = patches.permute(0, 1, 4, 7, 5, 8, 3, 2, 6, 9)
466
+
467
+ flatten_patches = patches.reshape(
468
+ batch_size,
469
+ grid_t * grid_h * grid_w,
470
+ channel * temporal_patch_size * patch_size * patch_size,
471
+ )
472
+
473
+ processed_images_grouped[shape] = flatten_patches
474
+ processed_grids[shape] = [[grid_t, grid_h, grid_w]] * batch_size
475
+
476
+ processed_images = reorder_images(
477
+ processed_images_grouped, grouped_images_index
478
+ )
479
+ processed_grids = reorder_images(processed_grids, grouped_images_index)
480
+
481
+ pixel_values = torch.cat(processed_images, dim=0)
482
+ image_grid_thw = torch.tensor(processed_grids)
483
+
484
+ return BatchFeature(
485
+ data={"pixel_values": pixel_values, "image_grid_thw": image_grid_thw},
486
+ tensor_type=return_tensors,
487
+ )
488
+
489
+ def preprocess(
490
+ self, images: ImageInput, **kwargs: Unpack[Glm5NextImageProcessorKwargs]
491
+ ) -> BatchFeature:
492
+ return super().preprocess(images, **kwargs)
493
+
494
+ def get_number_of_image_patches(
495
+ self, height: int, width: int, images_kwargs: dict | None = None
496
+ ) -> int:
497
+ """Number of image patches (pre-merge) for a given (height, width)."""
498
+ images_kwargs = images_kwargs or {}
499
+ patch_size = images_kwargs.get("patch_size", self.patch_size)
500
+ merge_size = images_kwargs.get("merge_size", self.merge_size)
501
+ patch_expand_factor = images_kwargs.get(
502
+ "patch_expand_factor", self.patch_expand_factor
503
+ )
504
+ min_pixels, max_pixels = _pixel_budget(
505
+ images_kwargs.get("min_image_tokens", self.min_image_tokens),
506
+ images_kwargs.get("max_image_tokens", self.max_image_tokens),
507
+ patch_size,
508
+ merge_size,
509
+ self.temporal_patch_size,
510
+ )
511
+ resized_height, resized_width = smart_resize(
512
+ t=self.temporal_patch_size,
513
+ h=height,
514
+ w=width,
515
+ t_factor=self.temporal_patch_size,
516
+ h_factor=patch_size * merge_size * patch_expand_factor,
517
+ w_factor=patch_size * merge_size * patch_expand_factor,
518
+ min_pixels=min_pixels,
519
+ max_pixels=max_pixels,
520
+ )
521
+ grid_h, grid_w = resized_height // patch_size, resized_width // patch_size
522
+ return grid_h * grid_w
523
+
524
+
525
+class Glm5NextVideoProcessorKwargs(VideosKwargs, total=False): # type: ignore[call-arg]
526
+ fps: list[float] | float
527
+ patch_size: int
528
+ temporal_patch_size: int
529
+ merge_size: int
530
+ patch_expand_factor: int
531
+ resize_mode: str | None
532
+ target_fps: float | None
533
+ max_frames: int | None
534
+ fps_interval: int | None
535
+ max_frame_count_dynamic: int | None
536
+ min_image_tokens: int | None
537
+ max_image_tokens: int | None
538
+
539
+
540
+class Glm5NextVideoProcessor(BaseVideoProcessor):
541
+ """Fast video processor for GLM-5.3-Flash.
542
+
543
+ Shares ``smart_resize`` / the pad-mode geometry / the patchify with the
544
+ image processor, and adds GLM-5.3-Flash frame sampling
545
+ (``glm_sample_frame_indices``: ``fps_interval`` semantics with a
546
+ temporal-patch-scaled greedy walk). Defaults mirror the checkpoint's
547
+ ``video_processor`` config.
548
+ """
549
+
550
+ resample = PILImageResampling.BICUBIC
551
+ size = {"longest_edge": 1} # unused: budgets come from the token bounds
552
+ image_mean = OPENAI_CLIP_MEAN
553
+ image_std = OPENAI_CLIP_STD
554
+ do_resize = True
555
+ do_rescale = True
556
+ do_normalize = True
557
+ do_convert_rgb = True
558
+ do_sample_frames = True
559
+ patch_size = 14
560
+ temporal_patch_size = 2
561
+ patch_expand_factor = 1
562
+ merge_size = 2
563
+ valid_kwargs = Glm5NextVideoProcessorKwargs
564
+ num_frames = 16
565
+ fps = 2
566
+ fps_interval = 2.0
567
+ max_frame_count_dynamic = 2048
568
+ resize_mode = "pad"
569
+ min_image_tokens = 16
570
+ max_image_tokens = 240000
571
+ model_input_names = ["pixel_values_videos", "video_grid_thw"]
572
+
573
+ def sample_frames(
574
+ self,
575
+ metadata: VideoMetadata,
576
+ fps: int | float | None = None,
577
+ **kwargs,
578
+ ) -> np.ndarray:
579
+ """Sample frame indices with GLM's fps-interval policy.
580
+
581
+ ``fps`` / ``target_fps``, ``max_frames`` and ``fps_interval`` /
582
+ ``max_frame_count_dynamic`` are the overrides described in
583
+ :func:`glm_sample_frame_indices`.
584
+ """
585
+ if metadata is None or getattr(metadata, "fps", None) is None:
586
+ raise ValueError(
587
+ "Asked to sample frames per second but no video metadata was "
588
+ "provided which is required when sampling in GLM-5.3-Flash. Please "
589
+ "pass in `VideoMetadata` object or set `do_sample_frames=False`."
590
+ )
591
+
592
+ target_fps = fps if fps is not None else kwargs.get("target_fps")
593
+ if target_fps is None:
594
+ target_fps = self.fps_interval
595
+ indices = glm_sample_frame_indices(
596
+ metadata.total_num_frames,
597
+ metadata.fps,
598
+ metadata.duration or 0,
599
+ target_fps=target_fps,
600
+ max_frame_count=kwargs.get("max_frames") or self.max_frame_count_dynamic,
601
+ temporal_patch_size=self.temporal_patch_size,
602
+ )
603
+ return np.array(indices)
604
+
605
+ def _preprocess(
606
+ self,
607
+ videos: list[torch.Tensor],
608
+ do_convert_rgb: bool = True,
609
+ do_resize: bool = True,
610
+ size: SizeDict | None = None,
611
+ resample: "PILImageResampling | int | None" = PILImageResampling.BICUBIC,
612
+ do_rescale: bool = True,
613
+ rescale_factor: float = 1 / 255.0,
614
+ do_normalize: bool = True,
615
+ image_mean: float | list[float] | None = None,
616
+ image_std: float | list[float] | None = None,
617
+ patch_size: int | None = None,
618
+ temporal_patch_size: int | None = None,
619
+ patch_expand_factor: int | None = None,
620
+ merge_size: int | None = None,
621
+ resize_mode: str | None = None,
622
+ min_image_tokens: int | None = None,
623
+ max_image_tokens: int | None = None,
624
+ return_tensors: str | TensorType | None = None,
625
+ **kwargs,
626
+ ) -> BatchFeature:
627
+ patch_expand_factor = self.patch_expand_factor
628
+ patch_size = patch_size if patch_size is not None else self.patch_size
629
+ temporal_patch_size = (
630
+ temporal_patch_size
631
+ if temporal_patch_size is not None
632
+ else self.temporal_patch_size
633
+ )
634
+ merge_size = merge_size if merge_size is not None else self.merge_size
635
+ resize_mode = resize_mode if resize_mode is not None else self.resize_mode
636
+ min_pixels, max_pixels = _pixel_budget(
637
+ min_image_tokens if min_image_tokens is not None else self.min_image_tokens,
638
+ max_image_tokens if max_image_tokens is not None else self.max_image_tokens,
639
+ patch_size,
640
+ merge_size,
641
+ temporal_patch_size,
642
+ )
643
+ grouped_videos, grouped_videos_index = group_videos_by_shape(videos)
644
+ resized_videos_grouped = {}
645
+ for shape, stacked_videos in grouped_videos.items():
646
+ if do_convert_rgb:
647
+ stacked_videos = self.convert_to_rgb(stacked_videos)
648
+ b, t_len, c, h, w = stacked_videos.shape
649
+ num_frames, height, width = t_len, h, w
650
+ if do_resize:
651
+ resized_height, resized_width = smart_resize(
652
+ t=num_frames,
653
+ h=height,
654
+ w=width,
655
+ t_factor=temporal_patch_size,
656
+ h_factor=patch_size * merge_size * patch_expand_factor,
657
+ w_factor=patch_size * merge_size * patch_expand_factor,
658
+ min_pixels=min_pixels,
659
+ max_pixels=max_pixels,
660
+ )
661
+ stacked_videos = stacked_videos.view(b * t_len, c, h, w)
662
+ stacked_videos = _resize_or_pad(
663
+ stacked_videos,
664
+ target_height=resized_height,
665
+ target_width=resized_width,
666
+ resize_mode=resize_mode,
667
+ resample=resample,
668
+ resize=self.resize,
669
+ allow_upscale=(num_frames * height * width < min_pixels),
670
+ )
671
+ stacked_videos = stacked_videos.view(
672
+ b, t_len, c, resized_height, resized_width
673
+ )
674
+ resized_videos_grouped[shape] = stacked_videos
675
+ resized_videos = reorder_videos(resized_videos_grouped, grouped_videos_index)
676
+
677
+ grouped_videos, grouped_videos_index = group_videos_by_shape(resized_videos)
678
+ processed_videos_grouped = {}
679
+ processed_grids = {}
680
+ for shape, stacked_videos in grouped_videos.items():
681
+ resized_height, resized_width = get_image_size(
682
+ stacked_videos[0], channel_dim=ChannelDimension.FIRST
683
+ )
684
+ stacked_videos = self.rescale_and_normalize(
685
+ stacked_videos,
686
+ do_rescale,
687
+ rescale_factor,
688
+ do_normalize,
689
+ image_mean,
690
+ image_std,
691
+ )
692
+ patches = stacked_videos
693
+
694
+ if pad := -patches.shape[1] % temporal_patch_size:
695
+ repeats = patches[:, -1:].expand(-1, pad, -1, -1, -1)
696
+ patches = torch.cat((patches, repeats), dim=1)
697
+ batch_size, grid_t, channel = patches.shape[:3]
698
+ grid_t = grid_t // temporal_patch_size
699
+ grid_h, grid_w = resized_height // patch_size, resized_width // patch_size
700
+
701
+ patches = patches.view(
702
+ batch_size,
703
+ grid_t,
704
+ temporal_patch_size,
705
+ channel,
706
+ grid_h // merge_size,
707
+ merge_size,
708
+ patch_size,
709
+ grid_w // merge_size,
710
+ merge_size,
711
+ patch_size,
712
+ )
713
+ patches = patches.permute(0, 1, 4, 7, 5, 8, 3, 2, 6, 9)
714
+ flatten_patches = patches.reshape(
715
+ batch_size,
716
+ grid_t * grid_h * grid_w,
717
+ channel * temporal_patch_size * patch_size * patch_size,
718
+ )
719
+
720
+ processed_videos_grouped[shape] = flatten_patches
721
+ processed_grids[shape] = [[grid_t, grid_h, grid_w]] * batch_size
722
+
723
+ processed_videos = reorder_videos(
724
+ processed_videos_grouped, grouped_videos_index
725
+ )
726
+ processed_grids = reorder_videos(processed_grids, grouped_videos_index)
727
+ pixel_values_videos = torch.cat(processed_videos, dim=0)
728
+ video_grid_thw = torch.tensor(processed_grids)
729
+ return BatchFeature(
730
+ data={
731
+ "pixel_values_videos": pixel_values_videos,
732
+ "video_grid_thw": video_grid_thw,
733
+ },
734
+ tensor_type=return_tensors,
735
+ )
736
+
737
+
738
+class Glm5NextProcessorKwargs(ProcessingKwargs, total=False): # type: ignore[call-arg]
739
+ images_kwargs: Glm5NextImageProcessorKwargs
740
+ videos_kwargs: Glm5NextVideoProcessorKwargs
741
+ _defaults = {
742
+ "text_kwargs": {
743
+ "padding": False,
744
+ "return_token_type_ids": False,
745
+ "return_mm_token_type_ids": False,
746
+ },
747
+ "videos_kwargs": {"return_metadata": True},
748
+ }
749
+
750
+
751
+class Glm5NextProcessor(ProcessorMixin):
752
+ """Wraps a GLM-5.3-Flash image processor, video processor and tokenizer.
753
+
754
+ Token expansion per image = ``prod(image_grid_thw) // merge_size**2``; video
755
+ frames are expanded with ``<|begin_of_image|>...<|end_of_image|>{ts} seconds``
756
+ structure (mrope timestamps).
757
+ """
758
+
759
+ attributes = ["image_processor", "tokenizer", "video_processor"]
760
+ image_processor_class = "AutoImageProcessor"
761
+ video_processor_class = "AutoVideoProcessor"
762
+ tokenizer_class = ("PreTrainedTokenizer", "PreTrainedTokenizerFast")
763
+
764
+ def __init__(
765
+ self,
766
+ image_processor=None,
767
+ tokenizer=None,
768
+ video_processor=None,
769
+ chat_template=None,
770
+ **kwargs,
771
+ ) -> None:
772
+ super().__init__(
773
+ image_processor, tokenizer, video_processor, chat_template=chat_template
774
+ )
775
+ self.image_token = (
776
+ "<|image|>"
777
+ if not hasattr(tokenizer, "image_token")
778
+ else tokenizer.image_token
779
+ )
780
+ self.video_token = (
781
+ "<|video|>"
782
+ if not hasattr(tokenizer, "video_token")
783
+ else tokenizer.video_token
784
+ )
785
+ self.image_token_id = (
786
+ tokenizer.image_token_id
787
+ if getattr(tokenizer, "image_token_id", None)
788
+ else tokenizer.convert_tokens_to_ids(self.image_token)
789
+ )
790
+ self.video_token_id = (
791
+ tokenizer.video_token_id
792
+ if getattr(tokenizer, "video_token_id", None)
793
+ else tokenizer.convert_tokens_to_ids(self.video_token)
794
+ )
795
+
796
+ @classmethod
797
+ def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
798
+ """Build the processor directly from the checkpoint config.
799
+
800
+ GLM-5.3-Flash stores nested image/video configs in
801
+ ``processor_config.json`` and declares a custom processor class. This
802
+ method reads those configs directly and caps only the video token budget.
803
+ """
804
+ from transformers import AutoTokenizer
805
+
806
+ model_path = pretrained_model_name_or_path
807
+ tokenizer = AutoTokenizer.from_pretrained(model_path, **kwargs)
808
+
809
+ def _cap_cfg(cfg: dict, *, is_video: bool) -> dict:
810
+ # Video keeps a serving token cap (the checkpoint's 240k-token
811
+ # budget would starve the KV cache at startup profiling); images
812
+ # follow the checkpoint budget verbatim so preprocessing matches
813
+ # the HF reference exactly.
814
+ if is_video and cfg.get("max_image_tokens") is not None:
815
+ cfg["max_image_tokens"] = min(
816
+ cfg["max_image_tokens"], _MAX_VIDEO_TOKENS
817
+ )
818
+ return cfg
819
+
820
+ ip_cfg = _cap_cfg(
821
+ dict(get_image_processor_config(model_path, **kwargs)), is_video=False
822
+ )
823
+ image_processor = Glm5NextImageProcessor(
824
+ **{k: v for k, v in ip_cfg.items() if k != "image_processor_type"}
825
+ )
826
+
827
+ processor_config = get_hf_file_to_dict(
828
+ "processor_config.json",
829
+ model_path,
830
+ revision=kwargs.get("revision", "main"),
831
+ )
832
+ if processor_config is None:
833
+ raise ValueError(f"Missing processor_config.json for {model_path}")
834
+ vp_cfg = _cap_cfg(dict(processor_config["video_processor"]), is_video=True)
835
+ video_processor = Glm5NextVideoProcessor(
836
+ **{k: v for k, v in vp_cfg.items() if k != "video_processor_type"}
837
+ )
838
+
839
+ return cls(
840
+ image_processor=image_processor,
841
+ tokenizer=tokenizer,
842
+ video_processor=video_processor,
843
+ )
844
+
845
+ def __call__(
846
+ self,
847
+ images: ImageInput | None = None,
848
+ text: TextInput
849
+ | PreTokenizedInput
850
+ | list[TextInput]
851
+ | list[PreTokenizedInput] = None,
852
+ videos: VideoInput | None = None,
853
+ **kwargs: Unpack[Glm5NextProcessorKwargs],
854
+ ) -> BatchFeature:
855
+ output_kwargs = self._merge_kwargs(
856
+ Glm5NextProcessorKwargs,
857
+ tokenizer_init_kwargs=self.tokenizer.init_kwargs,
858
+ **kwargs,
859
+ )
860
+ if images is not None:
861
+ image_inputs = self.image_processor(
862
+ images=images, **output_kwargs["images_kwargs"]
863
+ )
864
+ else:
865
+ image_inputs = {}
866
+
867
+ if videos is not None:
868
+ videos_inputs = self.video_processor(
869
+ videos=videos, **output_kwargs["videos_kwargs"]
870
+ )
871
+ if "return_metadata" not in kwargs:
872
+ videos_inputs.pop("video_metadata")
873
+ else:
874
+ videos_inputs = {}
875
+
876
+ if not isinstance(text, list):
877
+ text = [text]
878
+
879
+ # Prompt updates expand the unchanged image/video markers after this call.
880
+ return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None)
881
+ return_mm_token_type_ids = output_kwargs["text_kwargs"].pop(
882
+ "return_mm_token_type_ids", False
883
+ )
884
+ text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"])
885
+
886
+ if return_mm_token_type_ids:
887
+ array_ids = np.array(text_inputs["input_ids"])
888
+ mm_token_type_ids = np.zeros_like(text_inputs["input_ids"])
889
+ mm_token_type_ids[array_ids == self.image_token_id] = 1
890
+ text_inputs["mm_token_type_ids"] = mm_token_type_ids.tolist()
891
+ return BatchFeature(
892
+ data={**text_inputs, **image_inputs, **videos_inputs},
893
+ tensor_type=return_tensors,
894
+ )
895
+
896
+ def _get_num_multimodal_tokens(self, image_sizes=None, video_sizes=None, **kwargs):
897
+ vision_data = {}
898
+ if image_sizes is not None:
899
+ images_kwargs = Glm5NextProcessorKwargs._defaults.get("images_kwargs", {})
900
+ images_kwargs.update(kwargs)
901
+ merge_size = (
902
+ images_kwargs.get("merge_size", None) or self.image_processor.merge_size
903
+ )
904
+
905
+ num_image_patches = [
906
+ self.image_processor.get_number_of_image_patches(
907
+ *image_size, images_kwargs
908
+ )
909
+ for image_size in image_sizes
910
+ ]
911
+ num_image_tokens = [(n // merge_size**2) for n in num_image_patches]
912
+ vision_data.update(
913
+ {
914
+ "num_image_tokens": num_image_tokens,
915
+ "num_image_patches": num_image_patches,
916
+ }
917
+ )
918
+
919
+ if video_sizes is not None:
920
+ videos_kwargs = Glm5NextProcessorKwargs._defaults.get("videos_kwargs", {})
921
+ videos_kwargs.update(kwargs)
922
+ num_video_patches = [
923
+ self.video_processor.get_number_of_video_patches(
924
+ *video_size, videos_kwargs
925
+ )
926
+ for video_size in video_sizes
927
+ ]
928
+ num_video_tokens = [(n // merge_size**2) for n in num_video_patches]
929
+ vision_data["num_video_tokens"] = num_video_tokens
930
+
931
+ return MultiModalData(**vision_data)
932
+
933
+ def post_process_image_text_to_text(
934
+ self,
935
+ generated_outputs,
936
+ skip_special_tokens=True,
937
+ clean_up_tokenization_spaces=False,
938
+ **kwargs,
939
+ ):
940
+ return self.tokenizer.batch_decode(
941
+ generated_outputs,
942
+ skip_special_tokens=skip_special_tokens,
943
+ clean_up_tokenization_spaces=clean_up_tokenization_spaces,
944
+ **kwargs,
945
+ )
946
+
947
+
948
+__all__ = [
949
+ "Glm5NextImageProcessor",
950
+ "Glm5NextVideoProcessor",
951
+ "Glm5NextProcessor",
952
+ "smart_resize",
953
+]
vllm/utils/deep_gemm.py
+12
@@ -29,6 +29,11 @@ _DEEPGEMM_BLACKWELL_EXCLUDED_MODEL_TYPES: set[str] = {
29
"qwen3_5_moe_text",
30
}
31
32
+# KV page sizes (in cache entries) supported by the paged-MQA logits kernels
33
+# (fp8_fp4_paged_mqa_logits / get_paged_mqa_logits_metadata). Larger storage
34
+# blocks must be virtually split into one of these page sizes.
35
+PAGED_MQA_PAGE_SIZES = (32, 64)
36
+
37
38
def should_auto_disable_deep_gemm(model_type: str | None) -> bool:
39
"""Check if DeepGemm should be auto-disabled for this model on Blackwell.
@@ -648,6 +653,13 @@ def fp8_fp4_paged_mqa_logits(
653
_lazy_init()
654
if _fp8_fp4_paged_mqa_logits_impl is None:
655
return _missing()
656
+ # DeepGEMM asserts block_tables.stride(-1)==1. A trailing size-1 dim
657
+ # (e.g. block_table shape [B,1] for short seqs under a large block_size)
658
+ # can be a transposed view where .contiguous() is a no-op (torch treats the
659
+ # size-1 dim's stride as irrelevant) yet stride(-1)!=1, failing the kernel.
660
+ # clone to contiguous format to force stride(-1)==1.
661
+ if block_tables.dim() >= 2 and block_tables.stride(-1) != 1:
662
+ block_tables = block_tables.clone(memory_format=torch.contiguous_format)
663
kwargs = {} if indices is None else {"indices": indices}
664
return _fp8_fp4_paged_mqa_logits_impl(
665
q,
vllm/utils/flashinfer.py
+26
@@ -292,6 +292,32 @@ def has_flashinfer_moe() -> bool:
292
)
293
294
295
+@functools.cache
296
+def has_flashinfer_sm90_nope_mla() -> bool:
297
+ """FlashInfer SM90 NoPE MLA (FP8 KV with in-kernel dequant, kpe=0).
298
+
299
+ Feature-detected via the ``ckv_scale_arr`` run() kwarg introduced with
300
+ the SM90 NoPE support (FlashInfer >= 0.6.18), so dev builds carry the
301
+ gate without a version parse.
302
+ """
303
+ if not has_flashinfer():
304
+ return False
305
+ try:
306
+ import inspect
307
+
308
+ from flashinfer.mla import BatchMLAPagedAttentionWrapper
309
+ except ImportError:
310
+ return False
311
+ try:
312
+ params = inspect.signature(BatchMLAPagedAttentionWrapper.run).parameters
313
+ except (TypeError, ValueError):
314
+ return False
315
+ return (
316
+ "ckv_scale_arr" in params
317
+ and params["ckv_scale_arr"].kind is inspect.Parameter.KEYWORD_ONLY
318
+ )
319
+
320
+
321
@functools.cache
322
def has_flashinfer_sparse_mla_sm120() -> bool:
323
"""Return ``True`` if FlashInfer sparse MLA decode support is available."""
vllm/v1/attention/backend.py
+4
@@ -610,6 +610,10 @@ class AttentionMetadataBuilder(ABC, Generic[M]):
610
self.layer_names = layer_names
611
self.vllm_config = vllm_config
612
self.device = device
613
+ self.kernel_block_size: int | None = None
614
+
615
+ def set_kernel_block_size(self, kernel_block_size: int) -> None:
616
+ self.kernel_block_size = kernel_block_size
617
618
@classmethod
619
def get_cudagraph_support(
vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py
+54
-6
@@ -54,7 +54,10 @@ class _FlashInferMLASparseBackendBase(AttentionBackend):
54
55
@classmethod
56
def get_supported_head_sizes(cls) -> list[int]:
57
- return [576]
57
+ # 576 = 512 NoPE + 64 RoPE (with-rope layout); 512 = 512 NoPE only
58
+ # (no-rope layout, qk_rope_head_dim == 0). Both share D_V = 512 and are
59
+ # served by the TRTLLM-GEN sparse MLA kernel.
60
+ return [512, 576]
61
62
@classmethod
63
def is_mla(cls) -> bool:
@@ -117,8 +120,20 @@ class FlashInferMLASparseTRTLLMBackend(_FlashInferMLASparseBackendBase):
120
# FlashInfer MLA sparse SM10 kernel requires qk_nope_head_dim in [128, 192].
121
if vllm_config.model_config is not None:
122
hf_text_config = vllm_config.model_config.hf_text_config
120
- qk_nope_head_dim = getattr(hf_text_config, "qk_nope_head_dim", 1)
121
- if qk_nope_head_dim not in [128, 192]:
123
+ qk_nope_head_dim = hf_text_config.qk_nope_head_dim
124
+ qk_rope_head_dim = hf_text_config.qk_rope_head_dim
125
+ kv_lora_rank = hf_text_config.kv_lora_rank
126
+ if qk_rope_head_dim == 0:
127
+ # Native no-rope MLA: FlashInfer only ships the one shape
128
+ # (nope_mla_dimensions in flashinfer.mla._core).
129
+ if qk_nope_head_dim != 256 or kv_lora_rank != 512:
130
+ return (
131
+ "FlashInfer native no-rope MLA requires "
132
+ "qk_nope_head_dim=256 and kv_lora_rank=512, but got "
133
+ f"qk_nope_head_dim={qk_nope_head_dim}, "
134
+ f"kv_lora_rank={kv_lora_rank}"
135
+ )
136
+ elif qk_nope_head_dim not in [128, 192]:
137
return (
138
"FlashInfer MLA Sparse kernel requires qk_nope_head_dim "
139
f"in [128, 192], but got {qk_nope_head_dim}"
@@ -358,6 +373,10 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):
373
self.bmm1_scale: float | None = None
374
self.bmm2_scale: float | None = None
375
376
+ # Native no-rope MLA additionally requires a per-query-token active
377
+ # top-k length tensor.
378
+ self.is_nope_mla = self.qk_rope_head_dim == 0
379
+
380
# fp8 query quantization is required when using fp8 kv_cache,
381
# as the TRTLLM-GEN sparse MLA kernel requires matching dtypes
382
# for query and kv_cache (mixed bf16+fp8 is not supported).
@@ -429,6 +448,31 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):
448
block_tables = topk_indices_physical.unsqueeze(1)
449
seq_lens_arg = seq_lens
450
451
+ # page_table width = topk buffer width, which kpool widens past
452
+ # index_topk (topk_tokens) and rounds up to a multiple of 128. The
453
+ # kernel treats sparse_mla_top_k as the page-table *capacity* and bounds
454
+ # the active per-query length by ``seq_lens`` (the compacted valid
455
+ # count), so the -1 padding slots past seq_lens are never attended to.
456
+ # Use the actual buffer width instead of the fixed topk_tokens, which
457
+ # mismatches the page_table when index_kpool > 1.
458
+ sparse_topk_capacity = topk_indices_physical.shape[1]
459
+
460
+ extra_kwargs: dict[str, torch.Tensor] = {}
461
+ empty_rows: torch.Tensor | None = None
462
+ if self.is_nope_mla:
463
+ # The native no-rope kernel takes the active top-k length per query
464
+ # token (``seq_lens`` here is already the compacted per-token valid
465
+ # count, int32) and rejects zero-length rows. Point empty rows at a
466
+ # single valid dummy slot with length 1 and zero their output after
467
+ # the launch. ``triton_convert_req_index_to_global_index`` packs the
468
+ # valid indices into a contiguous prefix, which is what the kernel
469
+ # requires of the page table.
470
+ empty_rows = seq_lens == 0
471
+ topk_indices_physical[:, 0] = topk_indices_physical[:, 0].masked_fill(
472
+ empty_rows, 0
473
+ )
474
+ extra_kwargs["sparse_mla_top_k_lens"] = seq_lens.clamp(min=1)
475
+
476
kernel_out = trtllm_batch_decode_with_kv_cache_mla(
477
query=query,
478
kv_cache=kv_c_and_k_pe_cache.unsqueeze(1),
@@ -438,11 +482,12 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):
482
qk_rope_head_dim=self.qk_rope_head_dim,
483
block_tables=block_tables,
484
seq_lens=seq_lens_arg,
441
- max_seq_len=attn_metadata.topk_tokens,
485
+ max_seq_len=sparse_topk_capacity,
486
bmm1_scale=self.bmm1_scale,
487
bmm2_scale=self.bmm2_scale,
444
- sparse_mla_top_k=attn_metadata.topk_tokens,
488
+ sparse_mla_top_k=sparse_topk_capacity,
489
return_lse=self.need_to_return_lse_for_decode,
490
+ **extra_kwargs,
491
)
492
if self.need_to_return_lse_for_decode:
493
assert isinstance(kernel_out, tuple)
@@ -455,9 +500,12 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):
500
out = o.view(-1, o.shape[-2], o.shape[-1])
501
if lse is not None:
502
lse = self._normalize_lse(lse, out.shape[0], out.shape[1])
503
+ if empty_rows is None and lse is not None:
504
empty_rows = (topk_indices_physical == -1).all(dim=-1)
505
+ if empty_rows is not None:
506
out.masked_fill_(empty_rows.view(-1, 1, 1), 0.0)
460
- lse.masked_fill_(empty_rows.view(-1, 1), float("-inf"))
507
+ if lse is not None:
508
+ lse.masked_fill_(empty_rows.view(-1, 1), float("-inf"))
509
return out, lse
510
511
@staticmethod
vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py
new
+478
@@ -0,0 +1,478 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""FlashInfer sparse MLA backend for SM90 (Hopper) NoPE models.
4
+
5
+Wraps FlashInfer's ``BatchMLAPagedAttentionWrapper`` (FA2/FA3 paths), which
6
+as of FlashInfer 0.6.18 supports ``head_dim_kpe=0`` (GLM-5.3-Flash NoPE MLA)
7
+and FP8 E4M3 KV caches on SM90 with in-kernel dequantization: the FP8 cache
8
+is read directly (half the bf16 HBM traffic) and converted to BF16 in shared
9
+memory, while queries stay BF16 (no query quantization).
10
+
11
+Sparsity rides the same trick the FA-based sparse backend uses: with
12
+``page_size=1`` the per-token top-k slot indices ARE the page table, so each
13
+query token becomes one varlen batch row whose ``kv_indices`` slice is its
14
+top-k row and whose ``kv_len`` is its valid count. Causality is already
15
+encoded by the indexer's selection, so ``causal=False``.
16
+
17
+CUDA-graph handling: ``plan()`` copies its inputs to host unconditionally,
18
+so it must stay outside graph capture. Each metadata builder owns a wrapper,
19
+reserved capture-stable device buffers, and the plan parameters. The wrapper
20
+bakes the per-row ``kv_len`` into its int schedule
21
+at plan() time — ``run()`` never reads the device-side buffer — so the
22
+metadata builder replans every step (outside capture) with exact host-side
23
+lengths derived from the batch's sequence lengths; a full-width schedule
24
+would send the kernel past each row's valid count into the -1 tail of the
25
+converted index buffer (illegal address). Per-step content (top-k slots)
26
+is written into the reserved buffers by kernels inside the captured
27
+forward, and captured runs read the refreshed plan buffers on replay.
28
+
29
+KV cache format: plain contiguous E4M3 ``[num_blocks, block_size, 512]``
30
+(uint8 storage) with a per-tensor ``k_scale``; BF16 caches also work. The
31
+per-token x 128-channel-group ``ckv_scale_arr`` layout is supported by the
32
+kernel but not wired yet (it needs a group-quantizing cache-write op).
33
+"""
34
+
35
+from dataclasses import dataclass
36
+from typing import Any, ClassVar
37
+
38
+import torch
39
+
40
+from vllm.config import VllmConfig
41
+from vllm.config.cache import CacheDType
42
+from vllm.model_executor.layers.attention.sparse_mla_attention import (
43
+ SparseMLACommonImpl,
44
+)
45
+from vllm.platforms.interface import DeviceCapability
46
+from vllm.utils.flashinfer import has_flashinfer_sm90_nope_mla
47
+from vllm.v1.attention.backend import (
48
+ AttentionBackend,
49
+ AttentionLayer,
50
+ CommonAttentionMetadata,
51
+ MLAAttentionImpl,
52
+ MultipleOf,
53
+)
54
+from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
55
+ FlashInferMLASparseMetadata,
56
+ FlashInferMLASparseMetadataBuilder,
57
+)
58
+from vllm.v1.attention.backends.mla.sparse_utils import (
59
+ triton_convert_req_index_to_global_index,
60
+)
61
+from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout
62
+
63
+_FP8_KV_DTYPES = ("fp8", "fp8_e4m3")
64
+_WORKSPACE_BYTES = 128 * 1024 * 1024
65
+
66
+
67
+class FlashInferMLASparseSM90Backend(AttentionBackend):
68
+ supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
69
+ supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
70
+ "auto",
71
+ "bfloat16",
72
+ "fp8",
73
+ "fp8_e4m3",
74
+ ]
75
+
76
+ @staticmethod
77
+ def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
78
+ return [MultipleOf(64)]
79
+
80
+ @staticmethod
81
+ def get_name() -> str:
82
+ return "FLASHINFER_MLA_SPARSE_SM90"
83
+
84
+ @staticmethod
85
+ def get_builder_cls() -> type["FlashInferMLASparseSM90Builder"]:
86
+ return FlashInferMLASparseSM90Builder
87
+
88
+ @staticmethod
89
+ def get_impl_cls() -> type[MLAAttentionImpl]:
90
+ return FlashInferMLASparseSM90Impl
91
+
92
+ @classmethod
93
+ def get_supported_head_sizes(cls) -> list[int]:
94
+ # 512 = ckv 512 + kpe 0 (NoPE); 576 = ckv 512 + kpe 64.
95
+ return [512, 576]
96
+
97
+ @classmethod
98
+ def is_mla(cls) -> bool:
99
+ return True
100
+
101
+ @classmethod
102
+ def is_sparse(cls) -> bool:
103
+ return True
104
+
105
+ @classmethod
106
+ def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
107
+ return capability.major == 9
108
+
109
+ @classmethod
110
+ def supports_combination(
111
+ cls,
112
+ head_size: int,
113
+ dtype: torch.dtype,
114
+ kv_cache_dtype: CacheDType | None,
115
+ block_size: int | None,
116
+ use_mla: bool,
117
+ has_sink: bool,
118
+ use_sparse: bool,
119
+ use_mm_prefix: bool,
120
+ device_capability: DeviceCapability,
121
+ ) -> str | None:
122
+ if not has_flashinfer_sm90_nope_mla():
123
+ return (
124
+ "FLASHINFER_MLA_SPARSE_SM90 requires FlashInfer with SM90 "
125
+ "MLA support (ckv_scale_arr in "
126
+ "BatchMLAPagedAttentionWrapper.run, FlashInfer >= 0.6.18)"
127
+ )
128
+ if not use_sparse:
129
+ return "FLASHINFER_MLA_SPARSE_SM90 requires sparse MLA"
130
+ from vllm.config import get_current_vllm_config
131
+
132
+ vllm_config = get_current_vllm_config()
133
+ if vllm_config.model_config is not None:
134
+ hf = vllm_config.model_config.hf_text_config
135
+ # The SM90 FA2/FA3 kernel covers ckv=512 with kpe in {0, 64}
136
+ # (NoPE models and DeepSeek-style rope MLA alike).
137
+ if hf.kv_lora_rank != 512:
138
+ return "FLASHINFER_MLA_SPARSE_SM90 requires kv_lora_rank=512"
139
+ if hf.qk_rope_head_dim not in (0, 64):
140
+ return "FLASHINFER_MLA_SPARSE_SM90 requires qk_rope_head_dim in (0, 64)"
141
+ return None
142
+
143
+ @staticmethod
144
+ def get_kv_cache_shape(
145
+ num_blocks: int,
146
+ block_size: int,
147
+ num_kv_heads: int,
148
+ head_size: int,
149
+ cache_dtype_str: str = "auto",
150
+ ) -> tuple[int, ...]:
151
+ return (num_blocks, block_size, head_size)
152
+
153
+ @classmethod
154
+ def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]:
155
+ return (KVCacheLayout.LBHNC,)
156
+
157
+
158
+class _SM90State:
159
+ """Builder-owned wrapper, capture-stable buffers, and plan parameters.
160
+
161
+ One instance serves every MLA layer in an attention group because the plan
162
+ depends only on the batch shape, not the layer.
163
+ """
164
+
165
+ def __init__(
166
+ self,
167
+ device: torch.device,
168
+ num_heads: int,
169
+ kv_dtype: torch.dtype,
170
+ max_tokens: int,
171
+ topk_width: int,
172
+ kv_lora_rank: int,
173
+ qk_rope_head_dim: int,
174
+ sm_scale: float,
175
+ ) -> None:
176
+ from flashinfer.mla import BatchMLAPagedAttentionWrapper
177
+
178
+ self.workspace = torch.empty(_WORKSPACE_BYTES, dtype=torch.uint8, device=device)
179
+ self.device = device
180
+ self.num_heads = num_heads
181
+ self.kv_dtype = kv_dtype
182
+ self.max_tokens = max_tokens
183
+ self.topk_width = topk_width
184
+ self.kv_lora_rank = kv_lora_rank
185
+ self.qk_rope_head_dim = qk_rope_head_dim
186
+ self.sm_scale = sm_scale
187
+ # User-reserved buffers: with use_cuda_graph=True plan() refreshes
188
+ # these in place, so run()'s captured kernels always read them.
189
+ self.kv_indices = torch.zeros(
190
+ max_tokens * topk_width, dtype=torch.int32, device=device
191
+ )
192
+ self.kv_len_arr = torch.full(
193
+ (max_tokens,), topk_width, dtype=torch.int32, device=device
194
+ )
195
+ self.wrapper = BatchMLAPagedAttentionWrapper(
196
+ self.workspace,
197
+ qo_indptr=torch.zeros(max_tokens + 1, dtype=torch.int32, device=device),
198
+ kv_indptr=torch.zeros(max_tokens + 1, dtype=torch.int32, device=device),
199
+ kv_indices=self.kv_indices,
200
+ kv_len_arr=self.kv_len_arr,
201
+ use_cuda_graph=True,
202
+ backend="fa3",
203
+ )
204
+ self._arange_cpu = torch.arange(self.max_tokens + 1, dtype=torch.int32)
205
+ self._qo_cpu = torch.empty(self.max_tokens + 1, dtype=torch.int32)
206
+ self._kv_cpu = torch.empty(self.max_tokens + 1, dtype=torch.int32)
207
+ self._lens_cpu = torch.full(
208
+ (self.max_tokens,), self.topk_width, dtype=torch.int32
209
+ )
210
+
211
+ def plan(self, num_tokens: int, kv_lens: torch.Tensor) -> None:
212
+ """Replan with exact per-row KV lengths (CPU int32, ``[num_tokens]``).
213
+
214
+ The wrapper bakes kv_len into its int schedule from host values;
215
+ run() never reads the device kv_len_arr buffer. Scheduling with the
216
+ full buffer width while kv_indices rows carry a -1 tail past each
217
+ row's valid count makes the kernel compute ``-1 * ckv_stride_page``
218
+ (illegal address), so the lengths must be exact at plan time and
219
+ replanned every step as contexts grow. Must run outside CUDA graph
220
+ capture: the in-place refreshed plan_info/indptr buffers are what
221
+ captured runs read.
222
+ """
223
+ if torch.cuda.is_current_stream_capturing():
224
+ raise RuntimeError(
225
+ "FlashInferMLASparseSM90 plan() called inside CUDA graph "
226
+ "capture; lengths must be planned host-side before capture."
227
+ )
228
+ # CPU staging buffers are filled in place: plan() runs per step
229
+ # (once per draft/verify metadata build), so per-call allocations and
230
+ # device round trips are on the hot path. Passing CPU tensors lets
231
+ # the wrapper's internal .to("cpu") no-op; its reserved-buffer
232
+ # copy_ then performs the single H2D transfer per tensor.
233
+ # use_cuda_graph=True makes the wrapper copy qo/kv indptr into its
234
+ # fixed (max_tokens+1)-sized buffers with exact-size copy_, so the
235
+ # indptr must always be full-size. Rows past num_tokens are padded
236
+ # empty (qo_indptr flat at num_tokens) — zero-query rows read no q
237
+ # and schedule no work. Padded rows keep the full width lens; the
238
+ # value is never dereferenced.
239
+ torch.clamp(self._arange_cpu, max=num_tokens, out=self._qo_cpu)
240
+ torch.mul(self._qo_cpu, self.topk_width, out=self._kv_cpu)
241
+ self._lens_cpu.fill_(self.topk_width)
242
+ self._lens_cpu[:num_tokens] = kv_lens.to(torch.int32)
243
+ self.wrapper.plan(
244
+ self._qo_cpu,
245
+ self._kv_cpu,
246
+ self.kv_indices,
247
+ self._lens_cpu,
248
+ self.num_heads,
249
+ self.kv_lora_rank, # head_dim_ckv
250
+ self.qk_rope_head_dim, # 0 (NoPE) or 64 (rope MLA)
251
+ 1, # page_size: top-k slots are the page table
252
+ False, # causal: encoded by the indexer's selection
253
+ self.sm_scale,
254
+ q_data_type=torch.bfloat16,
255
+ kv_data_type=self.kv_dtype,
256
+ )
257
+
258
+
259
+@dataclass
260
+class FlashInferMLASparseSM90Metadata(FlashInferMLASparseMetadata):
261
+ state: _SM90State | None = None
262
+
263
+
264
+class FlashInferMLASparseSM90Builder(FlashInferMLASparseMetadataBuilder):
265
+ """Reuse the common sparse metadata (req ids, topk buffer access)."""
266
+
267
+ metadata_cls = FlashInferMLASparseSM90Metadata
268
+
269
+ def __init__(
270
+ self,
271
+ kv_cache_spec: "AttentionSpec",
272
+ layer_names: list[str],
273
+ vllm_config: "VllmConfig",
274
+ device: torch.device,
275
+ ) -> None:
276
+ super().__init__(kv_cache_spec, layer_names, vllm_config, device)
277
+ attention_layer = vllm_config.compilation_config.static_forward_context[
278
+ layer_names[0]
279
+ ]
280
+ impl = attention_layer.impl
281
+ if not isinstance(impl, FlashInferMLASparseSM90Impl):
282
+ raise TypeError(
283
+ "FlashInferMLASparseSM90Builder requires an SM90 FlashInfer "
284
+ f"implementation, got {type(impl).__name__}."
285
+ )
286
+ topk_indices_buffer = impl.topk_indices_buffer
287
+ assert topk_indices_buffer is not None
288
+ self.state = _SM90State(
289
+ device,
290
+ impl.num_heads,
291
+ kv_cache_spec.dtype,
292
+ vllm_config.scheduler_config.max_num_batched_tokens,
293
+ topk_indices_buffer.shape[1],
294
+ kv_lora_rank=impl.kv_lora_rank,
295
+ qk_rope_head_dim=impl.qk_rope_head_dim,
296
+ sm_scale=impl.scale,
297
+ )
298
+ # seq_lens_cpu_upper_bound is optimistic on decode rows under async
299
+ # spec decode, so the fast sync-free path is only safe without it;
300
+ # under async scheduling the exact (device) positions are used at
301
+ # the cost of one D2H sync per metadata build.
302
+ self._async_scheduling = bool(vllm_config.scheduler_config.async_scheduling)
303
+ hf_config = vllm_config.model_config.hf_text_config
304
+ assert hf_config.index_topk is not None
305
+ self._index_topk = int(hf_config.index_topk)
306
+ self._index_kpool = int(kv_cache_spec.tokens_per_state)
307
+
308
+ def _kv_lens_host(self, cam: CommonAttentionMetadata) -> tuple[int, torch.Tensor]:
309
+ """Exact per-row KV lengths, host-side (the flashinfer wrapper bakes
310
+ them into its schedule at plan time; there is no device-side path).
311
+
312
+ A row for the j-th query token of request i attends
313
+ ``seq_lens[i] - q_len[i] + j + 1`` tokens. The indexer's selection
314
+ then bounds the valid count: contexts up to ``index_topk`` select
315
+ everything (valid == context); longer contexts keep the top
316
+ ``index_topk`` pool-expanded tokens plus the trailing incomplete
317
+ pool (valid == ``index_topk + context % index_kpool``). Both match
318
+ the count of non -1 entries the convert kernel produces.
319
+ """
320
+ num_reqs = cam.num_reqs
321
+ qsl = cam.query_start_loc_cpu[: num_reqs + 1]
322
+ num_rows = int(qsl[-1])
323
+ if num_rows == 0:
324
+ return 0, torch.zeros(0, dtype=torch.int32)
325
+ # Row context == position + 1. Without async scheduling the
326
+ # maintained host upper bound is exact, giving a sync-free path;
327
+ # under async scheduling it is optimistic on decode rows, so fall
328
+ # back to the exact device positions (one D2H sync per build). The
329
+ # seq_lens derivation equals position + 1 because positions are
330
+ # contiguous per request.
331
+ sl_host = cam.seq_lens_cpu_upper_bound
332
+ positions = cam.positions
333
+ if not self._async_scheduling and sl_host is not None:
334
+ seq_lens = sl_host[:num_reqs].to(torch.int32)
335
+ q_lens = qsl[1:] - qsl[:-1]
336
+ first_pos = seq_lens - q_lens
337
+ req_of_row = torch.repeat_interleave(
338
+ torch.arange(num_reqs, dtype=torch.int64), q_lens.to(torch.int64)
339
+ )
340
+ rows = torch.arange(num_rows, dtype=torch.int32)
341
+ ctx = (
342
+ first_pos.to(torch.int64)[req_of_row]
343
+ + rows.to(torch.int64)
344
+ - qsl.to(torch.int64)[req_of_row]
345
+ + 1
346
+ )
347
+ elif positions is not None and num_rows <= positions.shape[0]:
348
+ ctx = positions[:num_rows].cpu().to(torch.int64) + 1
349
+ else:
350
+ seq_lens = cam.seq_lens[:num_reqs].cpu().to(torch.int32)
351
+ q_lens = qsl[1:] - qsl[:-1]
352
+ first_pos = seq_lens - q_lens
353
+ req_of_row = torch.repeat_interleave(
354
+ torch.arange(num_reqs, dtype=torch.int64), q_lens.to(torch.int64)
355
+ )
356
+ rows = torch.arange(num_rows, dtype=torch.int32)
357
+ ctx = (
358
+ first_pos.to(torch.int64)[req_of_row]
359
+ + rows.to(torch.int64)
360
+ - qsl.to(torch.int64)[req_of_row]
361
+ + 1
362
+ )
363
+ topk = self._index_topk
364
+ kpool = max(self._index_kpool, 1)
365
+ lens = torch.where(ctx <= topk, ctx, topk + ctx % kpool)
366
+ return num_rows, lens.to(torch.int32)
367
+
368
+ def build(
369
+ self,
370
+ common_prefix_len: int,
371
+ common_attn_metadata: CommonAttentionMetadata,
372
+ fast_build: bool = False,
373
+ ) -> FlashInferMLASparseSM90Metadata:
374
+ metadata = super().build(common_prefix_len, common_attn_metadata, fast_build)
375
+ assert isinstance(metadata, FlashInferMLASparseSM90Metadata)
376
+ # Replan every step outside any CUDA graph capture with this step's
377
+ # exact per-row lengths; captured runs read the refreshed buffers.
378
+ num_rows, kv_lens = self._kv_lens_host(common_attn_metadata)
379
+ self.state.plan(num_rows, kv_lens)
380
+ metadata.state = self.state
381
+ return metadata
382
+
383
+
384
+class FlashInferMLASparseSM90Impl(SparseMLACommonImpl[FlashInferMLASparseSM90Metadata]):
385
+ def __init__(
386
+ self,
387
+ num_heads: int,
388
+ head_size: int,
389
+ scale: float,
390
+ num_kv_heads: int,
391
+ alibi_slopes: list[float] | None,
392
+ sliding_window: int | None,
393
+ kv_cache_dtype: str,
394
+ logits_soft_cap: float | None,
395
+ attn_type: str,
396
+ kv_sharing_target_layer_name: str | None,
397
+ topk_indices_buffer: torch.Tensor | None = None,
398
+ indexer: Any | None = None,
399
+ **mla_args: Any,
400
+ ) -> None:
401
+ if any([alibi_slopes, sliding_window, logits_soft_cap]):
402
+ raise NotImplementedError(
403
+ "FlashInferMLASparseSM90Impl does not support alibi, sliding "
404
+ "window, or logits soft cap."
405
+ )
406
+ super().__init__(
407
+ num_heads,
408
+ head_size,
409
+ scale,
410
+ num_kv_heads,
411
+ alibi_slopes,
412
+ sliding_window,
413
+ kv_cache_dtype,
414
+ logits_soft_cap,
415
+ attn_type,
416
+ kv_sharing_target_layer_name,
417
+ indexer=indexer,
418
+ topk_indices_buffer=topk_indices_buffer,
419
+ **mla_args,
420
+ )
421
+ assert self.topk_indices_buffer is not None
422
+ self.supports_quant_query_input = False
423
+ self.use_fp8_kv_cache = self.kv_cache_dtype in _FP8_KV_DTYPES
424
+
425
+ def forward_mqa(
426
+ self,
427
+ q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
428
+ kv_c_and_k_pe_cache: torch.Tensor,
429
+ attn_metadata: FlashInferMLASparseSM90Metadata,
430
+ layer: AttentionLayer,
431
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
432
+ if not isinstance(q, tuple):
433
+ raise NotImplementedError(
434
+ "FlashInferMLASparseSM90Impl expects split (q_nope, q_rope)."
435
+ )
436
+ q_nope, q_rope = q
437
+ num_tokens = q_rope.shape[0]
438
+ # NoPE models hand a zero-width rope tensor through; rope MLA hands
439
+ # the real 64-dim part. The kernel takes both as-is.
440
+ q_pe = q_rope.reshape(num_tokens, self.num_heads, self.qk_rope_head_dim)
441
+
442
+ assert self.topk_indices_buffer is not None
443
+ topk_indices = self.topk_indices_buffer[:num_tokens]
444
+ # return_valid_counts=True keeps the compacted-prefix layout: valid
445
+ # entries at [0, valid_count), -1 past it — exactly the prefix the
446
+ # planned per-row lengths address.
447
+ topk_slots, _ = triton_convert_req_index_to_global_index(
448
+ attn_metadata.req_id_per_token[:num_tokens],
449
+ attn_metadata.block_table,
450
+ topk_indices,
451
+ BLOCK_SIZE=attn_metadata.block_size,
452
+ NUM_TOPK_TOKENS=topk_indices.shape[1],
453
+ return_valid_counts=True,
454
+ )
455
+ state = attn_metadata.state
456
+ assert state is not None
457
+ # Refresh top-k rows in graph and clamp masked tails to a valid slot;
458
+ # per-row lengths are already baked into the host-side plan.
459
+ width = topk_slots.shape[1]
460
+ state.kv_indices[: num_tokens * width].copy_(
461
+ topk_slots.reshape(-1).clamp_(min=0).to(torch.int32)
462
+ )
463
+
464
+ flat = (
465
+ kv_c_and_k_pe_cache.view(torch.float8_e4m3fn)
466
+ if self.use_fp8_kv_cache
467
+ else kv_c_and_k_pe_cache
468
+ ).reshape(-1, 1, self.head_size)
469
+ ckv = flat[..., : self.kv_lora_rank]
470
+ kpe = flat[..., self.kv_lora_rank :]
471
+
472
+ scale_kwargs = (
473
+ {"ckv_scale": float(layer._k_scale_float or 1.0), "kpe_scale": 1.0}
474
+ if self.use_fp8_kv_cache
475
+ else {}
476
+ )
477
+ out = state.wrapper.run(q_nope, q_pe, ckv, kpe, **scale_kwargs)
478
+ return out, None
vllm/v1/attention/backends/mla/indexer.py
+177
-5
@@ -37,7 +37,12 @@ from vllm.v1.attention.backends.utils import (
37
get_dcp_local_seq_lens,
38
split_decodes_and_prefills,
39
)
40
-from vllm.v1.kv_cache_interface import KVCacheLayout, KVCacheSpec, MLAAttentionSpec
40
+from vllm.v1.kv_cache_interface import (
41
+ AttentionSpec,
42
+ KVCacheLayout,
43
+ KVCacheSpec,
44
+ MLAAttentionSpec,
45
+)
46
47
logger = init_logger(__name__)
48
@@ -252,6 +257,30 @@ class DeepseekV32IndexerBackend(AttentionBackend):
257
return DeepseekV32IndexerMetadataBuilder
258
259
260
+class KpoolTailBackend(DeepseekV32IndexerBackend):
261
+ """Storage-only backend for the GLM-5.3-Flash kpool tail cache."""
262
+
263
+ @classmethod
264
+ def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]:
265
+ return (KVCacheLayout.LBHNC,)
266
+
267
+ @staticmethod
268
+ def get_name() -> str:
269
+ return "KPOOL_TAIL"
270
+
271
+ @classmethod
272
+ def get_supported_head_sizes(cls) -> list[int]:
273
+ return []
274
+
275
+ @staticmethod
276
+ def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
277
+ return [MultipleOf(1)]
278
+
279
+ @staticmethod
280
+ def get_builder_cls() -> type["KpoolTailMetadataBuilder"]: # type: ignore[override]
281
+ return KpoolTailMetadataBuilder
282
+
283
+
284
class DeepseekV4IndexerBackend(DeepseekV32IndexerBackend):
285
@staticmethod
286
def get_name() -> str:
@@ -419,6 +448,9 @@ class BuildPrefillChunkMetadataKernel(
448
)
449
)
450
)
451
+ index_kpool = getattr(hf_config, "index_kpool", None)
452
+ if index_kpool and index_kpool > 1 and index_kpool not in compress_ratios:
453
+ compress_ratios = compress_ratios + (index_kpool,)
454
return self._trace_dispatch(self.dispatch)(
455
# Cover Triton's divisible, exact-one, and generic i32 classes.
456
query_slice_start=(0, 1, 2),
@@ -483,6 +515,7 @@ class BuildPrefillChunkMetadataKernel(
515
@dataclass
516
class DeepseekV32IndexerPrefillMetadata:
517
chunks: list[DeepseekV32IndexerPrefillChunkMetadata]
518
+ max_prefill_seq_len: int = -1
519
520
521
@dataclass
@@ -497,6 +530,9 @@ class DeepSeekV32IndexerDecodeMetadata:
530
requires_padding: bool
531
schedule_metadata: torch.Tensor
532
global_seq_lens: torch.Tensor | None = None
533
+ per_req_decode_lens: torch.Tensor | None = None
534
+ decode_is_uniform: bool = True
535
+ write_max_decode_len: int = 0
536
indices: torch.Tensor | None = None
537
538
@@ -519,6 +555,90 @@ class DeepseekV32IndexerMetadata:
555
prefill: DeepseekV32IndexerPrefillMetadata | None = None
556
557
558
+def compute_kpool_tail_slot_mapping(
559
+ slot_mapping: torch.Tensor,
560
+ block_table: torch.Tensor,
561
+ query_start_loc: torch.Tensor,
562
+ positions: torch.Tensor,
563
+ num_actual_tokens: int,
564
+ num_reqs: int,
565
+ kpool: int,
566
+ out: torch.Tensor | None = None,
567
+) -> torch.Tensor:
568
+ """Map every token to its request's one circular tail block."""
569
+ if out is None:
570
+ out = slot_mapping.clone()
571
+ else:
572
+ assert out.shape == slot_mapping.shape
573
+ out.copy_(slot_mapping)
574
+ if num_actual_tokens == 0:
575
+ return out
576
+ tokens = torch.arange(num_actual_tokens, device=slot_mapping.device)
577
+ req = torch.searchsorted(query_start_loc, tokens, right=True) - 1
578
+ req = req.clamp_(min=0, max=num_reqs - 1)
579
+ own_block = block_table[:num_reqs, 0].index_select(0, req).to(torch.int64)
580
+ pos = positions[:num_actual_tokens].to(torch.int64)
581
+ out[:num_actual_tokens] = own_block * kpool + torch.remainder(pos, kpool)
582
+ return out
583
+
584
+
585
+class KpoolTailMetadataBuilder(AttentionMetadataBuilder):
586
+ """Build only the circular slot mapping needed by the storage-only tail."""
587
+
588
+ _cudagraph_support = AttentionCGSupport.ALWAYS
589
+ supports_update_block_table = False
590
+ reorder_batch_threshold = None
591
+
592
+ def __init__(
593
+ self,
594
+ kv_cache_spec: AttentionSpec,
595
+ layer_names: list[str],
596
+ vllm_config: VllmConfig,
597
+ device: torch.device,
598
+ ):
599
+ super().__init__(kv_cache_spec, layer_names, vllm_config, device)
600
+ self.slot_mapping_buffer = torch.empty(
601
+ vllm_config.scheduler_config.max_num_batched_tokens,
602
+ dtype=torch.int64,
603
+ device=device,
604
+ )
605
+
606
+ def build(
607
+ self,
608
+ common_prefix_len: int,
609
+ common_attn_metadata: CommonAttentionMetadata,
610
+ fast_build: bool = False,
611
+ ) -> DeepseekV32IndexerMetadata:
612
+ num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
613
+ split_decodes_and_prefills(common_attn_metadata)
614
+ )
615
+ slot_mapping = common_attn_metadata.slot_mapping
616
+ positions = common_attn_metadata.positions
617
+ if positions is not None:
618
+ slot_mapping_buffer = self.slot_mapping_buffer[
619
+ : slot_mapping.numel()
620
+ ].view_as(slot_mapping)
621
+ slot_mapping = compute_kpool_tail_slot_mapping(
622
+ slot_mapping,
623
+ common_attn_metadata.block_table_tensor,
624
+ common_attn_metadata.query_start_loc,
625
+ positions,
626
+ common_attn_metadata.num_actual_tokens,
627
+ common_attn_metadata.num_reqs,
628
+ self.kv_cache_spec.block_size,
629
+ out=slot_mapping_buffer,
630
+ )
631
+ return DeepseekV32IndexerMetadata(
632
+ seq_lens=common_attn_metadata.seq_lens,
633
+ max_seq_len=common_attn_metadata.max_seq_len,
634
+ slot_mapping=slot_mapping,
635
+ num_decodes=num_decodes,
636
+ num_decode_tokens=num_decode_tokens,
637
+ num_prefills=num_prefills,
638
+ num_prefill_tokens=num_prefill_tokens,
639
+ )
640
+
641
+
642
def get_max_prefill_buffer_size(vllm_config: VllmConfig):
643
max_model_len = vllm_config.model_config.max_model_len
644
# NOTE(Chen): 40 is a magic number for controlling the prefill buffer size.
@@ -640,6 +760,11 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
760
dtype=torch.int32,
761
device=self.device,
762
)
763
+ self.per_req_decode_lens_buffer = torch.zeros(
764
+ (scheduler_config.max_num_batched_tokens,),
765
+ dtype=torch.int32,
766
+ device=self.device,
767
+ )
768
# Shared workspace for decode seq_lens. Native MTP views this as
769
# (B, max_decode_len) at runtime, keeping context_lens contiguous even
770
# when max_decode_len is smaller than next_n.
@@ -707,6 +832,8 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
832
dtype=torch.int32,
833
device=self.device,
834
)
835
+ self.indexer_decode_block_table_buffer: torch.Tensor | None = None
836
+ self._max_num_batched_tokens = scheduler_config.max_num_batched_tokens
837
838
def _dcp_localize_decode_seq_lens(
839
self,
@@ -921,7 +1048,16 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
1048
1049
compressed_slot_mapping = slot_mapping
1050
compressed_seq_lens = seq_lens
1051
+ indexer_block_table = block_table
1052
if self.compress_ratio > 1:
1053
+ kernel_block_size = self.kernel_block_size
1054
+ if (
1055
+ kernel_block_size is not None
1056
+ and self.kv_cache_spec.block_size != kernel_block_size
1057
+ and self.kv_cache_spec.block_size % kernel_block_size == 0
1058
+ ):
1059
+ factor = self.kv_cache_spec.block_size // kernel_block_size
1060
+ indexer_block_table = (block_table[:, ::factor] // factor).contiguous()
1061
padded_num_tokens = num_tokens
1062
if self.pcp_world_size > 1:
1063
padded_num_tokens = slot_mapping.shape[0] // self.pcp_world_size
@@ -929,7 +1065,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
1065
num_tokens,
1066
query_start_loc,
1067
seq_lens,
932
- block_table,
1068
+ indexer_block_table,
1069
self.kv_cache_spec.num_states,
1070
self.compress_ratio,
1071
out=self.compressed_slot_mapping_buffer,
@@ -979,7 +1115,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
1115
seq_lens,
1116
compressed_seq_lens,
1117
compressed_seq_lens_cpu,
982
- common_attn_metadata.block_table_tensor,
1118
+ indexer_block_table,
1119
self.compress_ratio,
1120
query_slice=query_slice,
1121
skip_kv_gather=query_slice.start > 0,
@@ -990,7 +1126,14 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
1126
# Skip when total_seq_lens is 0 (i.e., no compressed token).
1127
if metadata is not None:
1128
chunks.append(metadata)
993
- prefill_metadata = DeepseekV32IndexerPrefillMetadata(chunks)
1129
+ prefill_metadata = DeepseekV32IndexerPrefillMetadata(
1130
+ chunks,
1131
+ max_prefill_seq_len=(
1132
+ int(seq_lens_cpu[num_decodes:].max().item())
1133
+ if num_prefills > 0
1134
+ else 0
1135
+ ),
1136
+ )
1137
1138
decode_metadata = None
1139
if num_decodes > 0:
@@ -999,6 +1142,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
1142
out=self.decode_lens_buffer[:num_decodes],
1143
)
1144
decode_lens = self.decode_lens_buffer[:num_decodes]
1145
+ self.per_req_decode_lens_buffer[:num_decodes].copy_(decode_lens)
1146
decode_lens_cpu = torch.diff(
1147
common_attn_metadata.query_start_loc_cpu[: num_decodes + 1]
1148
)
@@ -1018,6 +1162,8 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
1162
block_table = common_attn_metadata.block_table_tensor[:num_decodes, ...]
1163
1164
max_decode_len = int(decode_lens_cpu.max().item())
1165
+ min_decode_len = int(decode_lens_cpu.min().item())
1166
+ write_is_uniform = min_decode_len == max_decode_len
1167
next_n = 1 + self.num_speculative_tokens
1168
# The kernel sees max_decode_len Q rows, not the configured next_n,
1169
# so legality is per-step: on SM90 a uniformly 3-deep batch has no
@@ -1064,7 +1210,30 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
1210
)
1211
)
1212
1067
- seq_lens_is_buffer_view = not use_native or next_n > 1
1213
+ if self.compress_ratio > 1:
1214
+ kernel_block_size = self.kernel_block_size
1215
+ if (
1216
+ kernel_block_size is not None
1217
+ and self.kv_cache_spec.block_size != kernel_block_size
1218
+ and self.kv_cache_spec.block_size % kernel_block_size == 0
1219
+ ):
1220
+ factor = self.kv_cache_spec.block_size // kernel_block_size
1221
+ compressed = block_table[:, ::factor] // factor
1222
+ rows, cols = compressed.shape
1223
+ if self.indexer_decode_block_table_buffer is None:
1224
+ self.indexer_decode_block_table_buffer = torch.zeros(
1225
+ (self._max_num_batched_tokens, cols),
1226
+ dtype=torch.int32,
1227
+ device=self.device,
1228
+ )
1229
+ self.indexer_decode_block_table_buffer[:rows, :cols].copy_(
1230
+ compressed
1231
+ )
1232
+ block_table = self.indexer_decode_block_table_buffer[:rows, :cols]
1233
+
1234
+ seq_lens_is_buffer_view = (use_native and next_n > 1) or (
1235
+ not use_native and max_decode_len > 1
1236
+ )
1237
1238
# DCP: localize the now-expanded per-token global bounds to this
1239
# rank's owned KV. Done here (after expansion) so each token's global
@@ -1113,6 +1282,9 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
1282
schedule_metadata=schedule_metadata,
1283
indices=decode_indices,
1284
global_seq_lens=global_seq_lens_for_decode,
1285
+ per_req_decode_lens=self.per_req_decode_lens_buffer[:num_decodes],
1286
+ decode_is_uniform=write_is_uniform,
1287
+ write_max_decode_len=max_decode_len,
1288
)
1289
1290
attn_metadata = DeepseekV32IndexerMetadata(
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+5
-1
@@ -881,7 +881,11 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
881
rows. A DCP rank holds ``1/dcp_world_size`` of the sequence, so the
882
request's block table is always wider than the pages a row can reach.
883
"""
884
- pages_per_block = self.kernel_block_size // self._segmented_page_size
884
+ kernel_block_size = self.kernel_block_size
885
+ page_size = self._segmented_page_size
886
+ assert kernel_block_size is not None
887
+ assert page_size is not None
888
+ pages_per_block = kernel_block_size // page_size
889
max_local_pages = row_block_table.shape[1]
890
max_local_blocks = cdiv(max_local_pages, pages_per_block)
891
assert max_local_blocks <= block_table.shape[1], (
vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
+151
-52
@@ -32,6 +32,9 @@ from vllm.v1.attention.backends.mla.rocm_aiter_mla import (
32
AiterMLAHelper,
33
)
34
from vllm.v1.attention.backends.utils import split_decodes_and_prefills
35
+from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
36
+ rocm_sparse_attn_prefill,
37
+)
38
from vllm.v1.kv_cache_interface import AttentionSpec
39
from vllm.v1.worker.workspace import current_workspace_manager
40
@@ -40,6 +43,51 @@ if TYPE_CHECKING:
43
logger = init_logger(__name__)
44
45
46
+def _use_rocm_sparse_triton(
47
+ *,
48
+ kv_cache_dtype: str,
49
+ head_size: int,
50
+ kv_lora_rank: int,
51
+ num_prefills: int,
52
+ num_decodes: int,
53
+ num_decode_tokens: int,
54
+ max_query_len: int,
55
+) -> bool:
56
+ """Select the rope-free BF16 path not supported by AITER sparse MLA."""
57
+ plain_decode = num_decode_tokens == num_decodes
58
+ return (
59
+ not kv_cache_dtype.startswith("fp8")
60
+ and head_size == kv_lora_rank
61
+ and plain_decode
62
+ and (num_prefills > 0 or (num_decodes > 0 and max_query_len == 1))
63
+ )
64
+
65
+
66
+def fit_kpool_indices_to_aiter(
67
+ token_indices: torch.Tensor, topk_tokens: int
68
+) -> torch.Tensor:
69
+ """Keep the live kpool tail while fitting AITER's fixed top-k width."""
70
+ if token_indices.shape[1] < topk_tokens:
71
+ raise ValueError("token_indices width must be at least topk_tokens")
72
+ if token_indices.shape[1] == topk_tokens:
73
+ return token_indices
74
+
75
+ history = token_indices[:, :topk_tokens]
76
+ tail = token_indices[:, topk_tokens:]
77
+ valid_history = (history >= 0).sum(dim=1)
78
+ valid_tail = (tail >= 0).sum(dim=1)
79
+ keep_history = torch.minimum(valid_history, topk_tokens - valid_tail)
80
+
81
+ columns = torch.arange(topk_tokens, device=token_indices.device).unsqueeze(0)
82
+ tail_offsets = columns - keep_history.unsqueeze(1)
83
+ tail_values = torch.gather(
84
+ tail, 1, tail_offsets.clamp(min=0, max=tail.shape[1] - 1)
85
+ )
86
+ output = torch.where(columns < keep_history.unsqueeze(1), history, tail_values)
87
+ valid_output = columns < (keep_history + valid_tail).unsqueeze(1)
88
+ return torch.where(valid_output, output, -1)
89
+
90
+
91
@triton.jit
92
def _convert_req_index_to_global_index_kernel(
93
req_id_ptr, # int32 [num_tokens]
@@ -359,6 +407,7 @@ class ROCMAiterMLASparseMetadataBuilder(
407
self.kv_cache_spec = kv_cache_spec
408
self.model_config = vllm_config.model_config
409
self.model_dtype = vllm_config.model_config.dtype
410
+ self.kv_cache_dtype = vllm_config.cache_config.cache_dtype
411
parallel_config = vllm_config.parallel_config
412
self.device = device
413
max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
@@ -550,53 +599,74 @@ class ROCMAiterMLASparseMetadataBuilder(
599
# context lengths (both clamped to topk_tokens, past which per-token KV
600
# length saturates) and num_heads; fingerprint those CPU-side and skip
601
# the launch when nothing changed.
553
- num_reqs = common_attn_metadata.num_reqs
554
- clamped_seq_lens = np.minimum(
555
- common_attn_metadata.seq_lens_cpu[:num_reqs].numpy(),
556
- self.topk_tokens,
557
- )
558
- clamped_context_lens = np.minimum(
559
- common_attn_metadata.seq_lens_cpu[:num_reqs].numpy() - seg_lengths,
560
- self.topk_tokens,
561
- )
562
- metadata_key = (
563
- num_tokens,
564
- int(common_attn_metadata.max_query_len),
565
- self._num_attention_heads,
566
- clamped_seq_lens.tobytes(),
567
- clamped_context_lens.tobytes(),
568
- seg_lengths.tobytes(),
602
+ head_size = self.mla_dims.kv_lora_rank + self.mla_dims.qk_rope_head_dim
603
+ use_triton_sparse = _use_rocm_sparse_triton(
604
+ kv_cache_dtype=self.kv_cache_dtype,
605
+ head_size=head_size,
606
+ kv_lora_rank=self.mla_dims.kv_lora_rank,
607
+ num_prefills=num_prefills,
608
+ num_decodes=num_decodes,
609
+ num_decode_tokens=num_decode_tokens,
610
+ max_query_len=common_attn_metadata.max_query_len,
611
)
570
- if metadata_key != self._prev_metadata_key:
571
- from aiter import get_mla_metadata_v1
572
-
573
- max_split_per_batch = self._sparse_decode_max_split(
574
- int(common_attn_metadata.max_seq_len)
612
+ work_meta_data = None
613
+ work_indptr = None
614
+ work_info_set = None
615
+ reduce_indptr = None
616
+ reduce_final_map = None
617
+ reduce_partial_map = None
618
+ if not use_triton_sparse:
619
+ num_reqs = common_attn_metadata.num_reqs
620
+ clamped_seq_lens = np.minimum(
621
+ common_attn_metadata.seq_lens_cpu[:num_reqs].numpy(),
622
+ self.topk_tokens,
623
+ )
624
+ clamped_context_lens = np.minimum(
625
+ common_attn_metadata.seq_lens_cpu[:num_reqs].numpy() - seg_lengths,
626
+ self.topk_tokens,
627
)
576
- get_mla_metadata_v1(
577
- qo_indptr,
578
- paged_kv_indptr,
579
- paged_kv_last_page_len,
628
+ metadata_key = (
629
+ num_tokens,
630
+ int(common_attn_metadata.max_query_len),
631
self._num_attention_heads,
581
- 1,
582
- True,
583
- self._mla_work_meta_data,
584
- self._mla_work_info_set,
585
- self._mla_work_indptr,
586
- self._mla_reduce_indptr,
587
- self._mla_reduce_final_map,
588
- self._mla_reduce_partial_map,
589
- page_size=1,
590
- kv_granularity=16,
591
- max_seqlen_qo=1,
592
- uni_seqlen_qo=1,
593
- fast_mode=True,
594
- max_split_per_batch=max_split_per_batch,
632
+ clamped_seq_lens.tobytes(),
633
+ clamped_context_lens.tobytes(),
634
+ seg_lengths.tobytes(),
635
)
596
- # The persistent metadata buffers are read by graph replay. Order
597
- # the async metadata write before the graph-captured decode kernel.
598
- torch.cuda.current_stream(self.device).synchronize()
599
- self._prev_metadata_key = metadata_key
636
+ if metadata_key != self._prev_metadata_key:
637
+ from aiter import get_mla_metadata_v1
638
+
639
+ max_split_per_batch = self._sparse_decode_max_split(
640
+ int(common_attn_metadata.max_seq_len)
641
+ )
642
+ get_mla_metadata_v1(
643
+ qo_indptr,
644
+ paged_kv_indptr,
645
+ paged_kv_last_page_len,
646
+ self._num_attention_heads,
647
+ 1,
648
+ True,
649
+ self._mla_work_meta_data,
650
+ self._mla_work_info_set,
651
+ self._mla_work_indptr,
652
+ self._mla_reduce_indptr,
653
+ self._mla_reduce_final_map,
654
+ self._mla_reduce_partial_map,
655
+ page_size=1,
656
+ kv_granularity=16,
657
+ max_seqlen_qo=1,
658
+ uni_seqlen_qo=1,
659
+ fast_mode=True,
660
+ max_split_per_batch=max_split_per_batch,
661
+ )
662
+ torch.cuda.current_stream(self.device).synchronize()
663
+ self._prev_metadata_key = metadata_key
664
+ work_meta_data = self._mla_work_meta_data
665
+ work_indptr = self._mla_work_indptr
666
+ work_info_set = self._mla_work_info_set
667
+ reduce_indptr = self._mla_reduce_indptr
668
+ reduce_final_map = self._mla_reduce_final_map
669
+ reduce_partial_map = self._mla_reduce_partial_map
670
671
metadata = ROCMAiterMLASparseMetadata(
672
num_reqs=common_attn_metadata.num_reqs,
@@ -617,12 +687,12 @@ class ROCMAiterMLASparseMetadataBuilder(
687
paged_kv_last_page_len=paged_kv_last_page_len,
688
paged_kv_indices=paged_kv_indices,
689
paged_kv_indptr=paged_kv_indptr,
620
- work_meta_data=self._mla_work_meta_data,
621
- work_indptr=self._mla_work_indptr,
622
- work_info_set=self._mla_work_info_set,
623
- reduce_indptr=self._mla_reduce_indptr,
624
- reduce_final_map=self._mla_reduce_final_map,
625
- reduce_partial_map=self._mla_reduce_partial_map,
690
+ work_meta_data=work_meta_data,
691
+ work_indptr=work_indptr,
692
+ work_info_set=work_info_set,
693
+ reduce_indptr=reduce_indptr,
694
+ reduce_final_map=reduce_final_map,
695
+ reduce_partial_map=reduce_partial_map,
696
)
697
return metadata
698
@@ -715,6 +785,31 @@ class ROCMAiterMLASparseImpl(MLAAttentionImpl[ROCMAiterMLASparseMetadata]):
785
device=q.device,
786
)
787
788
+ if _use_rocm_sparse_triton(
789
+ kv_cache_dtype=self.kv_cache_dtype,
790
+ head_size=q.shape[-1],
791
+ kv_lora_rank=self.kv_lora_rank,
792
+ num_prefills=attn_metadata.num_prefills,
793
+ num_decodes=attn_metadata.num_decodes,
794
+ num_decode_tokens=attn_metadata.num_decode_tokens,
795
+ max_query_len=attn_metadata.max_query_len,
796
+ ):
797
+ rocm_sparse_attn_prefill(
798
+ q=q,
799
+ kv=kv_c_and_k_pe_cache.view(-1, 1, q.shape[-1]),
800
+ indices=None,
801
+ topk_length=None,
802
+ scale=self.scale,
803
+ head_dim=q.shape[-1],
804
+ nope_head_dim=self.kv_lora_rank,
805
+ rope_head_dim=q.shape[-1] - self.kv_lora_rank,
806
+ attn_sink=None,
807
+ output=output,
808
+ ragged_indices=attn_metadata.paged_kv_indices,
809
+ ragged_indptr=attn_metadata.paged_kv_indptr,
810
+ )
811
+ return AiterMLAHelper.get_mla_unpadded_o(self.num_heads, output)
812
+
813
# Build kwargs and forward the persistent MLA metadata when it has
814
# been computed. The aiter mla_decode_fwd switches to its
815
# work-stealing persistent kernel path when work_meta_data is given.
@@ -766,13 +861,18 @@ class ROCMAiterMLASparseImpl(MLAAttentionImpl[ROCMAiterMLASparseMetadata]):
861
)
862
else:
863
q = self.q_concat_buffer[: ql_nope.shape[0]]
769
- ops.concat_mla_q(ql_nope, q_pe, q)
864
+ if q_pe.shape[-1] == 0:
865
+ q.copy_(ql_nope)
866
+ else:
867
+ ops.concat_mla_q(ql_nope, q_pe, q)
868
869
num_actual_toks = attn_metadata.num_actual_tokens
870
871
# Get topk indices
872
assert self.topk_indices_buffer is not None
775
- topk_indices = self.topk_indices_buffer[:num_actual_toks]
873
+ topk_indices = fit_kpool_indices_to_aiter(
874
+ self.topk_indices_buffer[:num_actual_toks], attn_metadata.topk_tokens
875
+ )
876
877
triton_convert_req_index_to_global_index(
878
attn_metadata.req_id_per_token,
@@ -795,5 +895,4 @@ class ROCMAiterMLASparseImpl(MLAAttentionImpl[ROCMAiterMLASparseMetadata]):
895
attn_out = self._forward_mla(
896
layer, mla_padded_q, kv_c_and_k_pe_cache, attn_metadata
897
)
798
-
898
return attn_out, None
vllm/v1/attention/backends/mla/sparse_utils.py
+11
-4
@@ -251,7 +251,7 @@ class ConvertReqIndexToGlobalIndexKernel(
251
dict(
252
HAS_PREFILL_WORKSPACE=False,
253
COUNT_VALID=True,
254
- COMPACT_TO_FRONT=False,
254
+ COMPACT_TO_FRONT=True,
255
DCP_SIZE=1,
256
DCP_RANK=0,
257
DCP_INTERLEAVE=1,
@@ -259,7 +259,7 @@ class ConvertReqIndexToGlobalIndexKernel(
259
dict(
260
HAS_PREFILL_WORKSPACE=True,
261
COUNT_VALID=True,
262
- COMPACT_TO_FRONT=False,
262
+ COMPACT_TO_FRONT=True,
263
DCP_SIZE=1,
264
DCP_RANK=0,
265
DCP_INTERLEAVE=1,
@@ -458,7 +458,14 @@ def triton_convert_req_index_to_global_index(
458
req_id_c = req_id.contiguous()
459
block_table_c = block_table.contiguous()
460
token_indices_c = token_indices.contiguous()
461
- out = torch.empty_like(token_indices_c)
461
+ # When return_valid_counts, the kernel scatters valid entries to a
462
+ # contiguous prefix [0, valid_count) and leaves the tail unwritten, so
463
+ # pre-fill -1 there. flash_mla_sparse_fwd then bounds attention to
464
+ # [:topk_length] == exactly the valid set (no dropped tokens).
465
+ if return_valid_counts:
466
+ out = torch.full_like(token_indices_c, -1)
467
+ else:
468
+ out = torch.empty_like(token_indices_c)
469
470
valid_counts: torch.Tensor | None = None
471
if return_valid_counts:
@@ -491,7 +498,7 @@ def triton_convert_req_index_to_global_index(
498
NUM_TOPK_TOKENS=NUM_TOPK_TOKENS,
499
HAS_PREFILL_WORKSPACE=HAS_PREFILL_WORKSPACE,
500
COUNT_VALID=return_valid_counts,
494
- COMPACT_TO_FRONT=False,
501
+ COMPACT_TO_FRONT=return_valid_counts,
502
# DCP disabled (no-op de-interleave)
503
DCP_SIZE=1,
504
DCP_RANK=0,
vllm/v1/attention/backends/registry.py
+4
@@ -77,6 +77,10 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
77
"vllm.v1.attention.backends.mla.flashinfer_mla_sparse."
78
"FlashInferMLASparseSM120Backend"
79
)
80
+ FLASHINFER_MLA_SPARSE_SM90 = (
81
+ "vllm.v1.attention.backends.mla.flashinfer_mla_sparse_sm90."
82
+ "FlashInferMLASparseSM90Backend"
83
+ )
84
TRITON_MLA = "vllm.v1.attention.backends.mla.triton_mla.TritonMLABackend"
85
CUTLASS_MLA = "vllm.v1.attention.backends.mla.cutlass_mla.CutlassMLABackend"
86
FLASHMLA = "vllm.v1.attention.backends.mla.flashmla.FlashMLABackend"
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+29
-5
@@ -909,6 +909,7 @@ def rocm_aiter_sparse_attn_indexer(
909
has_decode = layer_attn_metadata.num_decodes > 0
910
has_prefill = layer_attn_metadata.num_prefills > 0
911
num_decode_tokens = layer_attn_metadata.num_decode_tokens
912
+ topk_indices_buffer[: hidden_states.shape[0]] = -1
913
914
# during speculative decoding, k may be padded to the CUDA graph batch
915
# size while slot_mapping only covers actual tokens.
@@ -1251,15 +1252,31 @@ _DSV4_SPARSE_NOPE_DIM = 448
1252
_DSV4_SPARSE_ROPE_DIM = 64
1253
1254
1254
-def _validate_dsv4_sparse_dims(
1255
+def _validate_sparse_dims(
1256
head_dim: int,
1257
nope_head_dim: int,
1258
rope_head_dim: int,
1259
op_name: str,
1260
) -> None:
1261
+ assert head_dim > 0, f"{op_name} expected a positive head_dim, got {head_dim}"
1262
+ assert nope_head_dim > 0, (
1263
+ f"{op_name} expected a positive NoPE dimension, got {nope_head_dim}"
1264
+ )
1265
+ assert rope_head_dim >= 0, (
1266
+ f"{op_name} expected a non-negative RoPE dimension, got {rope_head_dim}"
1267
+ )
1268
assert head_dim == nope_head_dim + rope_head_dim, (
1269
f"{op_name} expected head_dim={nope_head_dim + rope_head_dim}, got {head_dim}"
1270
)
1271
+
1272
+
1273
+def _validate_dsv4_sparse_dims(
1274
+ head_dim: int,
1275
+ nope_head_dim: int,
1276
+ rope_head_dim: int,
1277
+ op_name: str,
1278
+) -> None:
1279
+ _validate_sparse_dims(head_dim, nope_head_dim, rope_head_dim, op_name)
1280
assert (
1281
nope_head_dim == _DSV4_SPARSE_NOPE_DIM
1282
and rope_head_dim == _DSV4_SPARSE_ROPE_DIM
@@ -1351,6 +1368,12 @@ def _as_int32_contiguous_1d(x: torch.Tensor) -> torch.Tensor:
1368
return x.to(torch.int32).contiguous()
1369
1370
1371
+@triton.jit
1372
+def _sparse_kv_row_offset(slot, stride):
1373
+ # A global token slot fits in int32, but its byte/element offset may not.
1374
+ return slot.to(tl.int64) * stride
1375
+
1376
+
1377
@triton.jit
1378
def _sparse_attn_prefill_ragged_kernel(
1379
q_ptr,
@@ -1414,7 +1437,7 @@ def _sparse_attn_prefill_ragged_kernel(
1437
1438
kv = tl.load(
1439
kv_ptr
1417
- + safe_slot[:, None] * kv_stride_n
1440
+ + _sparse_kv_row_offset(safe_slot[:, None], kv_stride_n)
1441
+ dim_offsets[None, :] * kv_stride_d,
1442
mask=valid[:, None] & dim_mask[None, :],
1443
other=0.0,
@@ -2573,7 +2596,7 @@ def _rocm_sparse_attn_prefill_ragged_triton(
2596
assert indptr.numel() == num_queries + 1, (
2597
f"expected indptr shape [{num_queries + 1}], got {indptr.shape}"
2598
)
2576
- _validate_dsv4_sparse_dims(
2599
+ _validate_sparse_dims(
2600
head_dim,
2601
nope_head_dim,
2602
rope_head_dim,
@@ -3125,7 +3148,7 @@ def _rocm_sparse_attn_decode_triton(
3148
def rocm_sparse_attn_prefill(
3149
q: torch.Tensor,
3150
kv: torch.Tensor,
3128
- indices: torch.Tensor,
3151
+ indices: torch.Tensor | None,
3152
topk_length: torch.Tensor | None,
3153
scale: float,
3154
head_dim: int,
@@ -3139,7 +3162,7 @@ def rocm_sparse_attn_prefill(
3162
assert kv.ndim == 3 and kv.shape[1] == 1, (
3163
f"ROCm Triton sparse prefill expects kv=[skv,1,d], got {kv.shape}"
3164
)
3142
- _validate_dsv4_sparse_dims(
3165
+ _validate_sparse_dims(
3166
head_dim,
3167
nope_head_dim,
3168
rope_head_dim,
@@ -3157,6 +3180,7 @@ def rocm_sparse_attn_prefill(
3180
rope_head_dim=rope_head_dim,
3181
)
3182
else:
3183
+ assert indices is not None
3184
indices_2d = indices.reshape(indices.shape[0], -1)
3185
output_chunk = _rocm_sparse_attn_prefill_triton(
3186
q=q,
vllm/v1/core/kv_cache_coordinator.py
+8
@@ -599,6 +599,9 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
599
# can be a multiple of hash_block_size.
600
self.hash_block_size = hash_block_size
601
self.dcp_world_size = dcp_world_size
602
+ # Only groups that participate in prefix caching must satisfy the
603
+ # divisibility constraint; groups that opt out (e.g. GLM-5.3-Flash kpool
604
+ # tail, block_size=kpool) are scratch buffers and excluded.
605
group_block_sizes = [
606
manager.block_size
607
for manager, group in zip(
@@ -679,6 +682,11 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
682
"""
683
self.attention_groups: list[SpecGroup] = []
684
for i, g in enumerate(self.kv_cache_config.kv_cache_groups):
685
+ # Skip groups that opt out of prefix caching (e.g. GLM-5.3-Flash
686
+ # kpool tail): their blocks are per-request scratch, never
687
+ # shareable, so they must not participate in hit lookup (their
688
+ # manager-level hooks already no-op). Their slot in the per-group
689
+ # hit tuple stays empty.
690
if not g.kv_cache_spec.prefix_cacheable:
691
continue
692
manager_cls = self.single_type_managers[i].__class__
vllm/v1/core/kv_cache_utils.py
+295
-1
@@ -10,7 +10,7 @@ from collections import defaultdict
10
from collections.abc import Callable, Iterable, Iterator, Sequence
11
from dataclasses import dataclass, replace
12
from functools import partial
13
-from typing import Any, NamedTuple, NewType, TypeAlias, overload
13
+from typing import Any, NamedTuple, NewType, TypeAlias, cast, overload
14
15
from vllm import envs
16
from vllm.config import VllmConfig
@@ -24,6 +24,7 @@ from vllm.v1.kv_cache_interface import (
24
ChunkedLocalAttentionSpec,
25
FullAttentionSpec,
26
HiddenStateCacheSpec,
27
+ KpoolTailSpec,
28
KVCacheConfig,
29
KVCacheGroupSpec,
30
KVCacheLayout,
@@ -1132,6 +1133,205 @@ def _get_kv_cache_groups_uniform_type(
1133
return [KVCacheGroupSpec(list(spec.kv_cache_specs.keys()), spec)]
1134
1135
1136
+def _pp_balanced_mamba_group_count(
1137
+ vllm_config: VllmConfig,
1138
+ mamba_layer_names: list[str],
1139
+ mla_layer_names: list[str],
1140
+) -> int | None:
1141
+ """Return a Mamba group count whose PP projections fit the MLA slots."""
1142
+ num_groups = cdiv(len(mamba_layer_names), len(mla_layer_names))
1143
+ pp_size = vllm_config.parallel_config.pipeline_parallel_size
1144
+ if pp_size == 1:
1145
+ return num_groups
1146
+
1147
+ from vllm.distributed.utils import get_pp_indices
1148
+ from vllm.model_executor.models.utils import extract_layer_index
1149
+
1150
+ total_layers = vllm_config.model_config.get_total_num_hidden_layers()
1151
+ mamba_indices = [extract_layer_index(name) for name in mamba_layer_names]
1152
+ mla_indices = [extract_layer_index(name) for name in mla_layer_names]
1153
+ for rank in range(pp_size):
1154
+ start, end = get_pp_indices(total_layers, rank, pp_size)
1155
+ num_mamba = sum(start <= index < end for index in mamba_indices)
1156
+ num_mla = sum(start <= index < end for index in mla_indices)
1157
+ if not num_mamba:
1158
+ continue
1159
+ if not num_mla:
1160
+ return None
1161
+ num_groups = max(num_groups, cdiv(num_mamba, num_mla))
1162
+ return num_groups
1163
+
1164
+
1165
+def _get_kv_cache_groups_glm5_next(
1166
+ vllm_config: VllmConfig,
1167
+ kv_cache_spec: dict[str, KVCacheSpec],
1168
+) -> list[KVCacheGroupSpec] | None:
1169
+ """Build GLM-5.3-Flash groups with Mamba/MLA and tail/indexer aliasing."""
1170
+ mamba_specs = {
1171
+ name: spec
1172
+ for name, spec in kv_cache_spec.items()
1173
+ if isinstance(spec, MambaSpec)
1174
+ }
1175
+ tail_specs = {
1176
+ name: spec
1177
+ for name, spec in kv_cache_spec.items()
1178
+ if isinstance(spec, KpoolTailSpec)
1179
+ }
1180
+ attn_specs = {
1181
+ name: spec
1182
+ for name, spec in kv_cache_spec.items()
1183
+ if not isinstance(spec, (MambaSpec, KpoolTailSpec))
1184
+ }
1185
+ if not mamba_specs or not all(
1186
+ type(spec) is MLAAttentionSpec for spec in attn_specs.values()
1187
+ ):
1188
+ return None
1189
+
1190
+ mla_specs = cast(dict[str, MLAAttentionSpec], attn_specs)
1191
+ idx_pages = {
1192
+ spec.page_size_bytes for spec in mla_specs.values() if spec.tokens_per_state > 1
1193
+ }
1194
+ if not idx_pages:
1195
+ return None
1196
+
1197
+ assert all(spec.page_size_padded is None for spec in mla_specs.values())
1198
+ assert len(idx_pages) == 1
1199
+ mla_names = [name for name, spec in mla_specs.items() if spec.tokens_per_state == 1]
1200
+ mla_pages = {mla_specs[name].page_size_bytes for name in mla_names}
1201
+ assert len(mla_pages) == 1
1202
+ mla_page = mla_pages.pop()
1203
+ uniform_spec = UniformTypeKVCacheSpecs.from_specs(attn_specs)
1204
+ assert uniform_spec is not None
1205
+
1206
+ tail_group: KVCacheGroupSpec | None = None
1207
+ if tail_specs:
1208
+ idx_page = next(iter(idx_pages))
1209
+ padded_tail_specs: dict[str, KVCacheSpec] = {
1210
+ name: replace(spec, page_size_padded=idx_page)
1211
+ for name, spec in tail_specs.items()
1212
+ }
1213
+ tail_uniform = UniformTypeKVCacheSpecs.from_specs(padded_tail_specs)
1214
+ assert tail_uniform is not None
1215
+ tail_group = KVCacheGroupSpec(list(padded_tail_specs), tail_uniform)
1216
+
1217
+ any_mamba = next(iter(mamba_specs.values()))
1218
+ assert all(spec == any_mamba for spec in mamba_specs.values())
1219
+ if any_mamba.real_page_size_bytes > mla_page:
1220
+ raise ValueError(
1221
+ f"the mamba state page ({any_mamba.real_page_size_bytes} bytes) "
1222
+ f"does not fit the MLA page ({mla_page} bytes); increase tensor "
1223
+ "parallelism or use a wider KV cache dtype"
1224
+ )
1225
+ padded_specs: dict[str, KVCacheSpec] = {
1226
+ name: replace(any_mamba, page_size_padded=mla_page) for name in mamba_specs
1227
+ }
1228
+ num_groups = _pp_balanced_mamba_group_count(
1229
+ vllm_config, list(mamba_specs), mla_names
1230
+ )
1231
+ if num_groups is None:
1232
+ raise ValueError(
1233
+ "a pipeline stage has mamba layers but no MLA layer to share "
1234
+ "slots with; realign the stage boundaries (VLLM_PP_LAYER_PARTITION)"
1235
+ )
1236
+ mamba_grouped_names: list[list[str]] = [[] for _ in range(num_groups)]
1237
+ for index, name in enumerate(mamba_specs):
1238
+ mamba_grouped_names[index % num_groups].append(name)
1239
+
1240
+ return (
1241
+ [KVCacheGroupSpec(list(attn_specs), uniform_spec)]
1242
+ + ([tail_group] if tail_group is not None else [])
1243
+ + create_kv_cache_group_specs(padded_specs, mamba_grouped_names)
1244
+ )
1245
+
1246
+
1247
+def _glm5_next_tensor_layout(
1248
+ kv_cache_groups: list[KVCacheGroupSpec],
1249
+) -> (
1250
+ tuple[
1251
+ KVCacheGroupSpec,
1252
+ list[KVCacheGroupSpec],
1253
+ list[str],
1254
+ list[str],
1255
+ int,
1256
+ int,
1257
+ list[str],
1258
+ int,
1259
+ ]
1260
+ | None
1261
+):
1262
+ """Recognize the GLM-5.3-Flash grouping after optional PP projection."""
1263
+ uniform_groups = [
1264
+ group
1265
+ for group in kv_cache_groups
1266
+ if isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs)
1267
+ ]
1268
+ mamba_groups = [
1269
+ group for group in kv_cache_groups if isinstance(group.kv_cache_spec, MambaSpec)
1270
+ ]
1271
+ attn_group: KVCacheGroupSpec | None = None
1272
+ tail_group: KVCacheGroupSpec | None = None
1273
+ for group in uniform_groups:
1274
+ inner = cast(UniformTypeKVCacheSpecs, group.kv_cache_spec).kv_cache_specs
1275
+ if all(type(spec) is MLAAttentionSpec for spec in inner.values()):
1276
+ attn_group = group
1277
+ elif all(isinstance(spec, KpoolTailSpec) for spec in inner.values()):
1278
+ tail_group = group
1279
+ if attn_group is None or not mamba_groups:
1280
+ return None
1281
+ if len(uniform_groups) + len(mamba_groups) != len(kv_cache_groups):
1282
+ return None
1283
+
1284
+ attn_uniform = cast(UniformTypeKVCacheSpecs, attn_group.kv_cache_spec)
1285
+ mla_inner = cast(dict[str, MLAAttentionSpec], attn_uniform.kv_cache_specs)
1286
+ if not all(
1287
+ type(spec) is MLAAttentionSpec and spec.page_size_padded is None
1288
+ for spec in mla_inner.values()
1289
+ ):
1290
+ return None
1291
+ mla_names = [
1292
+ name for name in attn_group.layer_names if mla_inner[name].tokens_per_state == 1
1293
+ ]
1294
+ idx_names = [
1295
+ name for name in attn_group.layer_names if mla_inner[name].tokens_per_state > 1
1296
+ ]
1297
+ mla_pages = {mla_inner[name].page_size_bytes for name in mla_names}
1298
+ idx_pages = {mla_inner[name].page_size_bytes for name in idx_names}
1299
+ if len(mla_pages) != 1 or len(idx_pages) != 1:
1300
+ return None
1301
+ mla_page = mla_pages.pop()
1302
+ idx_page = idx_pages.pop()
1303
+ if any(group.kv_cache_spec.page_size_bytes != mla_page for group in mamba_groups):
1304
+ return None
1305
+
1306
+ tail_names: list[str] = []
1307
+ tail_page = 0
1308
+ if tail_group is not None:
1309
+ tail_names = list(tail_group.layer_names)
1310
+ tail_inner = cast(
1311
+ UniformTypeKVCacheSpecs, tail_group.kv_cache_spec
1312
+ ).kv_cache_specs
1313
+ tail_pages = {
1314
+ cast(KpoolTailSpec, spec).unpadded_page_size_bytes
1315
+ for spec in tail_inner.values()
1316
+ }
1317
+ if len(tail_pages) != 1 or len(tail_names) != len(idx_names):
1318
+ return None
1319
+ tail_page = tail_pages.pop()
1320
+ if tail_page > idx_page:
1321
+ return None
1322
+
1323
+ return (
1324
+ attn_group,
1325
+ mamba_groups,
1326
+ mla_names,
1327
+ idx_names,
1328
+ mla_page,
1329
+ idx_page,
1330
+ tail_names,
1331
+ tail_page,
1332
+ )
1333
+
1334
+
1335
def unify_kv_cache_spec_page_size(
1336
kv_cache_spec: dict[str, KVCacheSpec],
1337
) -> dict[str, KVCacheSpec]:
@@ -1359,6 +1559,10 @@ def _get_kv_cache_bytes_per_block(
1559
kv_cache_groups: list[KVCacheGroupSpec],
1560
) -> int:
1561
"""Return the largest cache group's bytes per block."""
1562
+ if (glm5_layout := _glm5_next_tensor_layout(kv_cache_groups)) is not None:
1563
+ _, _, mla_names, idx_names, mla_page, idx_page, _, _ = glm5_layout
1564
+ return len(mla_names) * mla_page + len(idx_names) * idx_page
1565
+
1566
bytes_per_block = max(
1567
sum(
1568
_get_per_layer_spec(group, layer_name).page_size_bytes
@@ -1432,6 +1636,69 @@ def get_kv_cache_config_from_groups(
1636
),
1637
)
1638
1639
+ if (glm5_layout := _glm5_next_tensor_layout(kv_cache_groups)) is not None:
1640
+ (
1641
+ attn_group,
1642
+ mamba_groups,
1643
+ mla_names,
1644
+ idx_names,
1645
+ mla_page,
1646
+ idx_page,
1647
+ tail_names,
1648
+ _,
1649
+ ) = glm5_layout
1650
+ bytes_per_block = len(mla_names) * mla_page + len(idx_names) * idx_page
1651
+ num_blocks = may_override_num_blocks(
1652
+ vllm_config, available_memory // bytes_per_block
1653
+ )
1654
+ size = bytes_per_block * num_blocks
1655
+ attn_specs = cast(
1656
+ UniformTypeKVCacheSpecs, attn_group.kv_cache_spec
1657
+ ).kv_cache_specs
1658
+
1659
+ kv_cache_tensors: list[KVCacheTensor] = []
1660
+
1661
+ def add_tensor(layer_name: str, spec: KVCacheSpec, offset: int) -> None:
1662
+ kv_cache_tensors.append(
1663
+ KVCacheTensor(
1664
+ size=size,
1665
+ layers=[layer_name],
1666
+ layer_stride=spec.page_size_bytes * num_blocks,
1667
+ block_stride=spec.page_size_bytes,
1668
+ offset=offset,
1669
+ )
1670
+ )
1671
+
1672
+ for index, mla_name in enumerate(mla_names):
1673
+ offset = index * mla_page * num_blocks
1674
+ add_tensor(mla_name, attn_specs[mla_name], offset)
1675
+ for group in mamba_groups:
1676
+ if index < len(group.layer_names):
1677
+ add_tensor(group.layer_names[index], group.kv_cache_spec, offset)
1678
+
1679
+ idx_base = len(mla_names) * mla_page * num_blocks
1680
+ for index, idx_name in enumerate(idx_names):
1681
+ offset = idx_base + index * idx_page * num_blocks
1682
+ add_tensor(idx_name, attn_specs[idx_name], offset)
1683
+ if tail_names:
1684
+ tail_name = tail_names[index]
1685
+ tail_group = next(
1686
+ group for group in kv_cache_groups if tail_name in group.layer_names
1687
+ )
1688
+ tail_specs = cast(
1689
+ UniformTypeKVCacheSpecs, tail_group.kv_cache_spec
1690
+ ).kv_cache_specs
1691
+ add_tensor(tail_name, tail_specs[tail_name], offset)
1692
+
1693
+ return KVCacheConfig(
1694
+ num_blocks=num_blocks,
1695
+ kv_cache_tensors=kv_cache_tensors,
1696
+ kv_cache_groups=kv_cache_groups,
1697
+ prefix_cache_retention_interval=(
1698
+ vllm_config.cache_config.prefix_cache_retention_interval
1699
+ ),
1700
+ )
1701
+
1702
layout = vllm_config.cache_config.get_resolved_kv_cache_layout()
1703
validate_kv_cache_layout(layout, kv_cache_groups)
1704
bytes_per_block = _get_kv_cache_bytes_per_block(kv_cache_groups)
@@ -1941,6 +2208,9 @@ def get_kv_cache_groups(
2208
# full attention, or all layers are sliding window attention with the
2209
# same window size). Put all layers into one group.
2210
return _get_kv_cache_groups_uniform_type(uniform_spec)
2211
+ elif glm5_groups := _get_kv_cache_groups_glm5_next(vllm_config, kv_cache_spec):
2212
+ return glm5_groups
2213
+
2214
# Hidden-state layers use their own block table and must not be absorbed
2215
# into a compatible attention bucket.
2216
hidden_specs = {
@@ -2065,6 +2335,30 @@ def _max_memory_usage_bytes_from_groups(
2335
if not kv_cache_groups:
2336
return 0
2337
2338
+ if (glm5_layout := _glm5_next_tensor_layout(kv_cache_groups)) is not None:
2339
+ (
2340
+ attn_group,
2341
+ mamba_groups,
2342
+ mla_names,
2343
+ idx_names,
2344
+ mla_page,
2345
+ idx_page,
2346
+ tail_names,
2347
+ _,
2348
+ ) = glm5_layout
2349
+ uniform_spec = cast(UniformTypeKVCacheSpecs, attn_group.kv_cache_spec)
2350
+ total_blocks = uniform_spec.max_memory_usage_pages(vllm_config)
2351
+ total_blocks += sum(
2352
+ cdiv(
2353
+ group.kv_cache_spec.max_memory_usage_bytes(vllm_config),
2354
+ group.kv_cache_spec.page_size_bytes,
2355
+ )
2356
+ for group in mamba_groups
2357
+ )
2358
+ if tail_names:
2359
+ total_blocks += 1
2360
+ return total_blocks * (len(mla_names) * mla_page + len(idx_names) * idx_page)
2361
+
2362
bytes_per_block = _pool_bytes_per_block(kv_cache_groups)
2363
total_blocks = 0
2364
for group in kv_cache_groups:
vllm/v1/core/single_type_kv_cache_manager.py
+13
@@ -6,6 +6,7 @@ from collections import defaultdict
6
from collections.abc import Sequence
7
from typing import ClassVar
8
9
+from vllm.logger import init_logger
10
from vllm.utils.math_utils import cdiv
11
from vllm.v1.core.block_pool import BlockPool
12
from vllm.v1.core.kv_cache_utils import (
@@ -22,6 +23,7 @@ from vllm.v1.kv_cache_interface import (
23
CrossAttentionSpec,
24
FullAttentionSpec,
25
HiddenStateCacheSpec,
26
+ KpoolTailSpec,
27
KVCacheSpec,
28
MambaSpec,
29
MLAAttentionSpec,
@@ -33,6 +35,8 @@ from vllm.v1.kv_cache_interface import (
35
from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
36
from vllm.v1.request import Request
37
38
+logger = init_logger(__name__)
39
+
40
41
class SingleTypeKVCacheManager(ABC):
42
"""
@@ -1210,6 +1214,10 @@ class CircularBufferManager(FullAttentionManager):
1214
return 0
1215
1216
1217
+class KpoolTailManager(CircularBufferManager):
1218
+ """One-block circular scratch manager for ``KpoolTailSpec``."""
1219
+
1220
+
1221
class ChunkedLocalAttentionManager(SingleTypeKVCacheManager):
1222
def __init__(self, kv_cache_spec: ChunkedLocalAttentionSpec, **kwargs) -> None:
1223
super().__init__(kv_cache_spec, **kwargs)
@@ -2085,6 +2093,11 @@ def register_all_kvcache_specs(vllm_config):
2093
SlidingWindowManager,
2094
uniform_type_base_spec=SlidingWindowMLASpec,
2095
)
2096
+ KVCacheSpecRegistry.register(
2097
+ KpoolTailSpec,
2098
+ KpoolTailManager,
2099
+ uniform_type_base_spec=KpoolTailSpec,
2100
+ )
2101
2102
KVCacheSpecRegistry.register(
2103
MambaSpec, MambaManager, uniform_type_base_spec=MambaSpec
vllm/v1/engine/core.py
+12
-1
@@ -333,8 +333,19 @@ class EngineCore:
333
vllm_config.cache_config.num_gpu_blocks = scheduler_kv_cache_config.num_blocks
334
kv_cache_groups = scheduler_kv_cache_config.kv_cache_groups
335
if kv_cache_groups:
336
+ # Exclude groups that opt out of prefix caching (e.g. GLM-5.3-Flash
337
+ # kpool tail, a 1-block/req scratch buffer with block_size=kpool):
338
+ # their small block_size would otherwise drag the global block_size
339
+ # below the real allocator block size and desync it from mamba.
340
+ participating = [
341
+ g.kv_cache_spec.block_size
342
+ for g in kv_cache_groups
343
+ if g.kv_cache_spec.prefix_cacheable
344
+ ]
345
vllm_config.cache_config.block_size = min(
337
- g.kv_cache_spec.block_size for g in kv_cache_groups
346
+ participating
347
+ if participating
348
+ else [g.kv_cache_spec.block_size for g in kv_cache_groups]
349
)
350
update_kv_cache_capacity(vllm_config, scheduler_kv_cache_config)
351
vllm/v1/kv_cache_interface.py
+34
-1
@@ -158,6 +158,7 @@ class KVCacheSpec:
158
159
@property
160
def prefix_cacheable(self) -> bool:
161
+ """Whether this spec's group participates in prefix caching."""
162
return True
163
164
@property
@@ -555,6 +556,8 @@ class MLAAttentionSpec(FullAttentionSpec):
556
# DeepseekV4 only fields. Non-DeepseekV4 MLA models leave these at defaults.
557
alignment: int | None = None # Default to None for no padding.
558
model_version: str | None = None
559
+ storage_block_size: int | None = None
560
+ """Token width used to view storage when it differs from the kernel block."""
561
# Marks draft groups that flatten a non-causal query block into decode rows.
562
non_causal_multi_token_decode: bool = False
563
# MLA stores a single latent vector per state; there is no separate V.
@@ -572,13 +575,16 @@ class MLAAttentionSpec(FullAttentionSpec):
575
cache_dtype_str_set = set(spec.cache_dtype_str for spec in specs)
576
tokens_per_state_set = set(spec.tokens_per_state for spec in specs)
577
model_version_set = set(spec.model_version for spec in specs)
578
+ storage_block_size_set = set(spec.storage_block_size for spec in specs)
579
assert (
580
len(cache_dtype_str_set) == 1
581
and len(tokens_per_state_set) == 1
582
and len(model_version_set) == 1
583
+ and len(storage_block_size_set) == 1
584
), (
585
"All attention layers in the same KV cache group must use the same "
581
- "quantization method, tokens per state, and model version."
586
+ "quantization method, tokens per state, model version, and storage "
587
+ "block size."
588
)
589
non_causal_mtd_set = {spec.non_causal_multi_token_decode for spec in specs}
590
assert len(non_causal_mtd_set) == 1, (
@@ -597,6 +603,7 @@ class MLAAttentionSpec(FullAttentionSpec):
603
cache_dtype_str=cache_dtype_str_set.pop(),
604
tokens_per_state=tokens_per_state_set.pop(),
605
model_version=model_version_set.pop(),
606
+ storage_block_size=storage_block_size_set.pop(),
607
non_causal_multi_token_decode=non_causal_mtd_set.pop(),
608
)
609
for spec in specs:
@@ -853,6 +860,28 @@ class SlidingWindowMLASpec(SlidingWindowSpec):
860
)
861
862
863
+@dataclass(frozen=True, kw_only=True)
864
+class KpoolTailSpec(SlidingWindowSpec):
865
+ """One-block circular scratch cache for a kpool indexer's raw tail."""
866
+
867
+ def max_admission_blocks_per_request(
868
+ self, max_in_flight_tokens: int, max_model_len: int
869
+ ) -> int:
870
+ return 1
871
+
872
+ def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int:
873
+ return 1
874
+
875
+ def is_uniform_with_collection(
876
+ self, kv_cache_specs: dict[str, KVCacheSpec]
877
+ ) -> bool:
878
+ return all(isinstance(spec, KpoolTailSpec) for spec in kv_cache_specs.values())
879
+
880
+ @property
881
+ def prefix_cacheable(self) -> bool:
882
+ return False
883
+
884
+
885
@dataclass(frozen=True)
886
class MambaSpec(KVCacheSpec):
887
shapes: tuple[tuple[int, ...], ...]
@@ -875,6 +904,10 @@ class MambaSpec(KVCacheSpec):
904
for (shape, dtype) in zip(self.shapes, self.dtypes)
905
)
906
907
+ @property
908
+ def real_page_size_bytes(self) -> int:
909
+ return self.state_content_size_bytes
910
+
911
@property
912
def page_size_bytes(self) -> int:
913
page_size = sum(
vllm/v1/spec_decode/llm_base_proposer.py
+1
@@ -1021,6 +1021,7 @@ class SpecDecodeBaseProposer:
1021
{
1022
"DeepSeekMTPModel",
1023
"DeepseekV32MTPModel",
1024
+ "Glm5NextMTPModel",
1025
"KimiK3MTPModel",
1026
}.intersection(architectures)
1027
)
vllm/v1/worker/gpu/model_runner.py
+1
@@ -696,6 +696,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
696
attn_groups_iter=(g for groups in self.attn_groups for g in groups),
697
kernel_block_sizes=self.kernel_block_sizes,
698
static_forward_context=self.compilation_config.static_forward_context,
699
+ num_blocks=self.kv_cache_config.num_blocks,
700
)
701
702
@torch.inference_mode()
vllm/v1/worker/gpu/model_states/mamba_hybrid.py
+1
@@ -322,6 +322,7 @@ class MambaHybridModelState(DefaultModelState):
322
kv_cache_config=kv_cache_config,
323
seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
324
dcp_local_seq_lens=input_batch.dcp_local_seq_lens,
325
+ positions=input_batch.positions,
326
model_specific_attn_metadata=mamba_attn_metadata,
327
for_cudagraph_capture=for_capture,
328
rswa_prefix_lens=input_batch.prompt_lens,
vllm/v1/worker/gpu_model_runner.py
+1
@@ -1160,6 +1160,7 @@ class GPUModelRunner(
1160
kernel_block_sizes=self._kernel_block_sizes,
1161
runner_only_attn_layers=self.runner_only_attn_layers,
1162
static_forward_context=self.compilation_config.static_forward_context,
1163
+ num_blocks=self.kv_cache_config.num_blocks,
1164
)
1165
1166
def _zero_block_ids(self, block_ids: list[int]) -> None:
vllm/v1/worker/utils.py
+26
-7
@@ -35,6 +35,7 @@ from vllm.v1.kv_cache_interface import (
35
KVCacheLayout,
36
KVCacheSpec,
37
MambaSpec,
38
+ MLAAttentionSpec,
39
UniformTypeKVCacheSpecs,
40
create_kv_cache_views,
41
)
@@ -115,6 +116,7 @@ class KVBlockZeroer:
116
attn_groups_iter: Iterable["AttentionGroup"],
117
kernel_block_sizes: list[int],
118
static_forward_context: dict[str, Any],
119
+ num_blocks: int,
120
runner_only_attn_layers: set[str] | None = None,
121
) -> None:
122
"""Precompute the absolute-address table for the Triton zeroing kernel.
@@ -158,8 +160,6 @@ class KVBlockZeroer:
160
continue
161
kernel_bs = kernel_block_sizes[group.kv_cache_group_id]
162
assert spec.block_size % kernel_bs == 0
161
- ratio = spec.block_size // kernel_bs
162
-
163
for layer_name in group.layer_names:
164
if layer_name in runner_only_attn_layers:
165
continue
@@ -168,6 +168,12 @@ class KVBlockZeroer:
168
continue
169
dp = kv.data_ptr()
170
171
+ assert kv.shape[0] % num_blocks == 0, (
172
+ f"{layer_name}: {kv.shape[0]} kernel blocks is not a "
173
+ f"multiple of {num_blocks} logical blocks"
174
+ )
175
+ ratio = kv.shape[0] // num_blocks
176
+
177
el = kv.element_size()
178
block_stride_bytes = kv.stride(0) * el
179
assert block_stride_bytes % 4 == 0
@@ -267,11 +273,19 @@ class AttentionGroup:
273
kernel_block_size: int | None = None,
274
num_metadata_builders: int = 1,
275
):
270
- kv_cache_spec_builder = (
271
- self.kv_cache_spec.copy_with_new_block_size(kernel_block_size)
272
- if kernel_block_size is not None
273
- else self.kv_cache_spec
274
- )
276
+ if kernel_block_size is None:
277
+ kv_cache_spec_builder = self.kv_cache_spec
278
+ elif (
279
+ isinstance(self.kv_cache_spec, MLAAttentionSpec)
280
+ and self.kv_cache_spec.storage_block_size is not None
281
+ ):
282
+ kv_cache_spec_builder = self.kv_cache_spec.copy_with_new_block_size(
283
+ self.kv_cache_spec.storage_block_size
284
+ )
285
+ else:
286
+ kv_cache_spec_builder = self.kv_cache_spec.copy_with_new_block_size(
287
+ kernel_block_size
288
+ )
289
builder_cls = self.backend.get_builder_cls()
290
builder_kwargs = {}
291
if builder_cls.requires_block_table_width:
@@ -291,6 +305,9 @@ class AttentionGroup:
305
)
306
for _ in range(num_metadata_builders)
307
]
308
+ if kernel_block_size is not None:
309
+ for builder in self.metadata_builders:
310
+ builder.set_kernel_block_size(kernel_block_size)
311
312
def get_metadata_builder(self, ubatch_id: int = 0) -> AttentionMetadataBuilder:
313
assert len(self.metadata_builders) > ubatch_id
@@ -427,6 +444,8 @@ def allocate_kv_cache(
444
kernel_block_size = None
445
if kernel_block_sizes is not None and group_id < len(kernel_block_sizes):
446
kernel_block_size = kernel_block_sizes[group_id]
447
+ if isinstance(spec, MLAAttentionSpec) and spec.storage_block_size is not None:
448
+ kernel_block_size = spec.storage_block_size
449
450
views = create_kv_cache_views(
451
buf,