@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 4fe9e6f6e5

[Refactor][GLM-5.3-Flash] Move sparse_attn_indexer_kpool into the model folder and split AMD/NVIDIA (#55358)

Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>

Jiangyun Zhu committed Sep 16, 2026 at 13:50 UTC 4fe9e6f6e5647a159c9e65e4f3bf33dec5d12e1c
6 files changed +1008 -398
tests/v1/attention/test_sparse_indexer_decode_seq_lens.py
+10 -6
@@ -20,10 +20,10 @@ kernel is documented to replicate.
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.
23 +# Bootstrap the glm5next package before importing kpool_compress directly: its
24 +# import runs glm5next/__init__, which pulls model -> attention -> the indexer
25 +# dispatcher (sparse_indexer) -> back into the nvidia/amd sparse_indexer
26 +# modules. 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
@@ -32,13 +32,17 @@ from vllm.models.glm5next.nvidia.ops.kpool_compress import ( # noqa: E402
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 (
35 +from vllm.models.glm5next.common.sparse_indexer import (
36 _decode_topk_seq_lens,
37 _fill_short_decode_causal_indices,
38 )
39 from vllm.platforms import current_platform
40
41 +if current_platform.is_rocm():
42 + from vllm.models.glm5next.amd import sparse_indexer as indexer_mod
43 +else:
44 + from vllm.models.glm5next.nvidia import sparse_indexer as indexer_mod
45 +
46 KPOOL = 4
47 TOPK_TOKENS = 16
48 SELECT_K = TOPK_TOKENS // KPOOL
vllm/models/glm5next/amd/sparse_indexer.py new
+729
@@ -0,0 +1,729 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Custom Sparse Attention Indexer layers."""
4 +
5 +import torch
6 +
7 +import vllm.envs as envs
8 +from vllm import _custom_ops # noqa: F401 # registers the torch.ops._C kernels
9 +from vllm._aiter_ops import rocm_aiter_ops
10 +from vllm.compilation.breakable_cudagraph import eager_break_during_capture
11 +from vllm.config import get_current_vllm_config_or_none
12 +from vllm.forward_context import get_forward_context
13 +from vllm.logger import init_logger
14 +from vllm.model_executor.custom_op import CustomOp
15 +from vllm.models.glm5next.amd.ops import kpool_compress as kpool_ops
16 +from vllm.models.glm5next.common.sparse_indexer import (
17 + RADIX_TOPK_WORKSPACE_SIZE,
18 + _build_decode_scatter_indices,
19 + _decode_topk_seq_lens,
20 + _fill_causal_indices,
21 + _fill_short_decode_causal_indices,
22 + _gather_workspace_shapes,
23 + _scatter_decode_tokens_by_request,
24 + kv_cache_as_quant_view,
25 +)
26 +from vllm.platforms import current_platform
27 +from vllm.utils.torch_utils import (
28 + LayerNameType,
29 + _encode_layer_name,
30 + _resolve_layer_name,
31 +)
32 +from vllm.v1.attention.backends.mla.indexer import (
33 + DeepseekV32IndexerMetadata,
34 +)
35 +from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton
36 +from vllm.v1.worker.workspace import current_workspace_manager
37 +
38 +logger = init_logger(__name__)
39 +
40 +# kpool write helper: form pools from the current token batch and compress them
41 +# into the index K cache via the fused Triton kernel.
42 +
43 +
44 +def _kpool_compress_insert(
45 + k: torch.Tensor,
46 + gate_score: torch.Tensor,
47 + ape: torch.Tensor,
48 + kv_cache: torch.Tensor,
49 + slot_mapping: torch.Tensor,
50 + kpool: int,
51 + head_dim: int,
52 + round_scale: bool,
53 +) -> None:
54 + """Pool ``kpool`` consecutive tokens into one fp8 K and write at pool slots.
55 +
56 + ``slot_mapping`` is pool-granular (compress_ratio == kpool on the spec):
57 + only the *last* token of each complete pool carries a valid (>=0) slot;
58 + intra-pool tokens are -1. Every position is treated as a pool-completion
59 + candidate and non-completions are masked off inside the kernel. Compacting
60 + the valid rows first costs two device syncs on the eager prefill path and
61 + buys nothing numerically. Assumes pool-aligned chunk starts.
62 + """
63 + n = slot_mapping.shape[0]
64 + # No pool can complete in a batch smaller than one pool; also keeps the
65 + # clamped gather indices below in bounds.
66 + if n < kpool:
67 + return
68 + pos = torch.arange(n, device=k.device)
69 + valid = slot_mapping >= 0
70 + # Drop pools whose start falls before the batch (leading padding); their
71 + # gate/k data is undefined anyway.
72 + write_mask = valid & (pos >= kpool - 1)
73 + offs = torch.arange(kpool, device=k.device)
74 + idx = (pos - (kpool - 1)).clamp_min(0)[:, None] + offs[None, :]
75 + kpool_ops.kpool_compress_and_write_cache(
76 + kv_cache,
77 + k[idx], # [n, kpool, head_dim]
78 + gate_score[idx],
79 + ape,
80 + slot_mapping.to(torch.int64),
81 + pool_size=kpool,
82 + head_dim=head_dim,
83 + write_mask=write_mask,
84 + round_scale=round_scale,
85 + write_cache=True,
86 + return_compressed=False,
87 + )
88 +
89 +
90 +@eager_break_during_capture
91 +def sparse_attn_indexer_kpool(
92 + hidden_states: torch.Tensor,
93 + k_cache_prefix: LayerNameType,
94 + kv_cache: torch.Tensor,
95 + q_quant: torch.Tensor,
96 + q_scale: torch.Tensor | None,
97 + k: torch.Tensor,
98 + weights: torch.Tensor,
99 + quant_block_size: int,
100 + scale_fmt: str | None,
101 + topk_tokens: int,
102 + head_dim: int,
103 + max_model_len: int,
104 + total_seq_lens: int,
105 + topk_indices_buffer: torch.Tensor,
106 + skip_k_cache_insert: bool,
107 + use_fp4_cache: bool = False,
108 + # kpool params (Plan-A: gate is consumed at write time and read back at
109 + # topk time to softmax-weight the pool).
110 + gate_score: torch.Tensor | None = None,
111 + compress_ape: torch.Tensor | None = None,
112 + index_kpool: int = 1,
113 + positions: torch.Tensor | None = None,
114 + # Paged tail cache (in-progress pool's raw K + gate score), replacing the
115 + # transient _DECODE_TAIL ring. tail_prefix resolves attn_metadata[tail_prefix]
116 + # for the tail group's token-granular slot_mapping. None on the dummy/profiling
117 + # path and when the tail cache is disabled.
118 + tail_kv_cache: torch.Tensor | None = None,
119 + tail_prefix: str | None = None,
120 +) -> torch.Tensor:
121 + # careful! this will be None in dummy run
122 + attn_metadata = get_forward_context().attn_metadata
123 + fp8_dtype = current_platform.fp8_dtype()
124 + k_cache_prefix = _resolve_layer_name(k_cache_prefix)
125 +
126 + # assert isinstance(attn_metadata, dict)
127 + if not isinstance(attn_metadata, dict):
128 + # Reserve workspace for indexer during profiling run
129 + values_spec, scales_spec = _gather_workspace_shapes(
130 + total_seq_lens, head_dim, fp8_dtype, use_fp4_cache
131 + )
132 + current_workspace_manager().get_simultaneous(
133 + values_spec,
134 + scales_spec,
135 + ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
136 + )
137 +
138 + # Reserve profiler-visible memory for the worst-case decode logits,
139 + # whose shape is [B * next_n, max_model_len]. This profiling branch
140 + # returns before invoking the logits kernel itself.
141 + cfg = get_current_vllm_config_or_none()
142 + worst_decode_tokens = 0
143 + if cfg is not None:
144 + sched = cfg.scheduler_config
145 + num_spec = (
146 + cfg.speculative_config.num_speculative_tokens
147 + if cfg.speculative_config is not None
148 + else 0
149 + )
150 + worst_decode_tokens = min(
151 + sched.max_num_seqs * (num_spec + 1),
152 + sched.max_num_batched_tokens,
153 + )
154 + # float32 logits -> 4 bytes/element; uint8 sentinel so elems == bytes.
155 + decode_logits_elems = worst_decode_tokens * max_model_len * 4
156 + prefill_cap_elems = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024
157 + max_logits_elems = max(decode_logits_elems, prefill_cap_elems)
158 + _ = torch.empty(
159 + max_logits_elems, dtype=torch.uint8, device=hidden_states.device
160 + )
161 +
162 + return topk_indices_buffer
163 + attn_metadata_narrowed = attn_metadata[k_cache_prefix]
164 + assert isinstance(attn_metadata_narrowed, DeepseekV32IndexerMetadata)
165 + slot_mapping = attn_metadata_narrowed.slot_mapping
166 + has_decode = attn_metadata_narrowed.num_decodes > 0
167 + has_prefill = attn_metadata_narrowed.num_prefills > 0
168 + num_decode_tokens = attn_metadata_narrowed.num_decode_tokens
169 +
170 + # q_scale is required iff the FP4 cache path is enabled; the FP8 path
171 + # folds the Q scale into `weights` inside fused_indexer_q_rope_quant.
172 + if use_fp4_cache:
173 + assert q_scale is not None, "use_fp4_cache=True requires q_scale"
174 + else:
175 + assert q_scale is None, "q_scale must be None when use_fp4_cache=False"
176 +
177 + # During speculative decoding, k may be padded to the CUDA graph batch
178 + # size while slot_mapping only covers actual tokens. Truncate k to avoid
179 + # out-of-bounds reads in the kernel.
180 + num_tokens = slot_mapping.shape[0]
181 + if k is not None:
182 + k = k[:num_tokens]
183 +
184 + if not skip_k_cache_insert:
185 + assert not use_fp4_cache, "Unfused FP4 Insert is not supported yet"
186 + if index_kpool > 1 and gate_score is not None and compress_ape is not None:
187 + # kpool prefill write: pool kpool consecutive prefill tokens via
188 + # softmax(gate+ape)-weighted sum -> Hadamard -> fp8 -> pool slots.
189 + # Decode tokens (the first num_decode_tokens in the batch) cannot be
190 + # pooled here — their pool's earlier tokens are not in this batch —
191 + # so they are deferred to the tail-buffer kernel in has_decode.
192 + # compress_ratio == index_kpool makes slot_mapping pool-granular.
193 + n_prefill = num_tokens - num_decode_tokens
194 + if n_prefill > 0:
195 + # decode tokens are batched first; prefill tokens follow.
196 + prefill_slice = slice(num_decode_tokens, num_tokens)
197 + _kpool_compress_insert(
198 + k[prefill_slice],
199 + gate_score[prefill_slice],
200 + compress_ape,
201 + kv_cache,
202 + slot_mapping[prefill_slice],
203 + index_kpool,
204 + head_dim,
205 + round_scale=(scale_fmt is not None),
206 + )
207 + # Persist each request's incomplete prefill pool so decode can
208 + # finish it, including after PD transfer. Tail slots use
209 + # ``pos % kpool`` within the request's tail block. Processing
210 + # only the batch's trailing tokens would miss all but the last
211 + # request in a multi-request prefill.
212 + if tail_kv_cache is not None and tail_prefix is not None:
213 + tail_meta = attn_metadata.get(_resolve_layer_name(tail_prefix))
214 + if tail_meta is not None:
215 + assert isinstance(tail_meta, DeepseekV32IndexerMetadata)
216 + kpool_ops.kpool_seed_tail_cache(
217 + tail_kv_cache,
218 + k[prefill_slice],
219 + gate_score[prefill_slice],
220 + tail_meta.slot_mapping[prefill_slice],
221 + index_kpool,
222 + head_dim,
223 + )
224 + else:
225 + # standard: per-token fp8 quant + scatter (all tokens).
226 + assert scale_fmt is not None
227 + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
228 + indexer_k_quant_and_cache_triton,
229 + )
230 +
231 + indexer_k_quant_and_cache_triton(
232 + k,
233 + kv_cache,
234 + slot_mapping,
235 + quant_block_size,
236 + scale_fmt,
237 + )
238 +
239 + topk_indices_buffer[: hidden_states.shape[0]] = -1
240 + if has_prefill:
241 + prefill_metadata = attn_metadata_narrowed.prefill
242 + assert prefill_metadata is not None
243 +
244 + # Short sequences select every pool, so skip sparse scoring and fill
245 + # the top-k buffer with all causal token indices. The index-K cache was
246 + # already written above.
247 + n_prefill_sf = num_tokens - num_decode_tokens
248 + # Host-side short-prefill predicate: max_prefill_seq_len is computed
249 + # in the metadata builder (exact for prefill rows) and equals
250 + # positions[prefill_slice].max() + 1, so this replaces a
251 + # positions.max().item() device sync per layer. -1 (unknown metadata)
252 + # falls back to the device-side check.
253 + if prefill_metadata.max_prefill_seq_len >= 0:
254 + short_prefill = (
255 + n_prefill_sf > 0
256 + and positions is not None
257 + and prefill_metadata.max_prefill_seq_len <= topk_tokens
258 + )
259 + else:
260 + short_prefill = (
261 + n_prefill_sf > 0
262 + and positions is not None
263 + and int(positions[num_decode_tokens:num_tokens].max().item()) + 1
264 + <= topk_tokens
265 + )
266 + if short_prefill:
267 + # short_prefill is only True when positions is not None (above),
268 + # but narrow explicitly for the indexer below.
269 + assert positions is not None
270 + _pos = positions[num_decode_tokens:num_tokens].to(torch.int32)
271 + _buf = topk_indices_buffer[num_decode_tokens:num_tokens]
272 + _fill_causal_indices(_buf, _pos)
273 +
274 + # Get the full shared workspace buffers once (will allocate on first use).
275 + # Layout switches between FP8 (head_dim bytes + 4-byte fp32 scale) and
276 + # MXFP4 (head_dim/2 bytes packed + head_dim/MXFP4_BLOCK_SIZE ue8m0
277 + # scales) based on use_fp4_cache.
278 + workspace_manager = current_workspace_manager()
279 + values_spec, scales_spec = _gather_workspace_shapes(
280 + total_seq_lens, head_dim, fp8_dtype, use_fp4_cache
281 + )
282 + k_quant_full, k_scale_full = workspace_manager.get_simultaneous(
283 + values_spec,
284 + scales_spec,
285 + )
286 + for chunk in prefill_metadata.chunks if not short_prefill else ():
287 + k_quant = k_quant_full[: chunk.total_seq_lens]
288 + k_scale = k_scale_full[: chunk.total_seq_lens]
289 +
290 + if not chunk.skip_kv_gather:
291 + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
292 + cp_gather_indexer_k_quant_cache_triton,
293 + )
294 +
295 + cp_gather_indexer_k_quant_cache_triton(
296 + kv_cache,
297 + k_quant,
298 + k_scale,
299 + chunk.block_table,
300 + chunk.cu_seq_lens,
301 + token_to_seq=chunk.token_to_seq,
302 + )
303 +
304 + q_slice = q_quant[chunk.token_start : chunk.token_end]
305 + q_scale_slice = (
306 + q_scale[chunk.token_start : chunk.token_end]
307 + if q_scale is not None
308 + else None
309 + )
310 + # DeepGEMM scalar-type tags (zero-copy): MXFP4 values → int8
311 + # (kPackedFP4), scales → int32 squeezed to 1-D kv_sf / 2-D q_sf.
312 + if use_fp4_cache:
313 + q_slice_cast = q_slice.view(torch.int8)
314 + k_quant_cast = k_quant.view(torch.int8)
315 + k_scale_cast = k_scale.view(torch.int32).squeeze(-1)
316 + else:
317 + q_slice_cast = q_slice
318 + k_quant_cast = k_quant
319 + k_scale_cast = k_scale.view(torch.float32).squeeze(-1)
320 + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
321 + rocm_fp8_mqa_logits,
322 + )
323 +
324 + assert q_scale_slice is None
325 + logits = rocm_fp8_mqa_logits(
326 + q_slice_cast,
327 + (k_quant_cast, k_scale_cast),
328 + weights[chunk.token_start : chunk.token_end],
329 + chunk.cu_seqlen_ks,
330 + chunk.cu_seqlen_ke,
331 + )
332 + num_rows = logits.shape[0]
333 +
334 + # kpool: logits are pool-granular (compress_ratio == index_kpool),
335 + # so topk selects pools. We pick topk_tokens // kpool pools then
336 + # expand each pool back to its kpool constituent tokens.
337 + select_k = topk_tokens // index_kpool if index_kpool > 1 else topk_tokens
338 + if index_kpool > 1:
339 + pool_topk = torch.empty(
340 + (num_rows, select_k), dtype=torch.int32, device=logits.device
341 + )
342 + topk_dst = pool_topk
343 + else:
344 + topk_dst = topk_indices_buffer[
345 + chunk.token_start : chunk.token_end, :topk_tokens
346 + ]
347 +
348 + torch.ops._C.top_k_per_row_prefill(
349 + logits,
350 + chunk.cu_seqlen_ks,
351 + chunk.cu_seqlen_ke,
352 + topk_dst,
353 + num_rows,
354 + logits.stride(0),
355 + logits.stride(1),
356 + select_k,
357 + )
358 +
359 + if index_kpool > 1:
360 + pool_ids = pool_topk.to(torch.int64)
361 + if positions is not None:
362 + # Fused expand-pools + append-tail into one Triton kernel
363 + # (replaces ~25 elementwise ops). seq_len is token-granular
364 + # (pos+1); the kernel derives pool_len internally.
365 + q_seq = (
366 + positions[chunk.token_start : chunk.token_end].to(torch.int32)
367 + + 1
368 + )
369 + expanded = kpool_ops.expand_pools_and_append_tail(
370 + pool_ids, q_seq, index_kpool
371 + )
372 + else:
373 + valid = pool_ids >= 0
374 + expanded = kpool_ops.expand_pools_to_tokens(
375 + pool_ids, valid, topk_tokens, index_kpool
376 + )
377 + topk_indices_buffer[
378 + chunk.token_start : chunk.token_end, : expanded.shape[-1]
379 + ] = expanded
380 +
381 + if has_decode:
382 + decode_metadata = attn_metadata_narrowed.decode
383 + assert decode_metadata is not None
384 + kv_cache_raw = kv_cache # raw [num_blocks, block_size, head_dim+4] for writes
385 + kv_cache = kv_cache_as_quant_view(kv_cache, head_dim, use_fp4_cache)
386 +
387 + # Update the tail before reading logits; completed pools are compressed
388 + # into the slot supplied by slot_mapping.
389 + # Spec verification groups tokens by request and preserves position
390 + # order so each token is stashed before the next completes its pool.
391 + # Positions must remain token-granular because the kernel derives the
392 + # pool phase and tail index from ``pos % kpool``.
393 + if (
394 + index_kpool > 1
395 + and gate_score is not None
396 + and compress_ape is not None
397 + and positions is not None
398 + and not skip_k_cache_insert
399 + ):
400 + num_requests = attn_metadata_narrowed.num_decodes
401 + # Kpool writes must recover the original request grouping after the
402 + # indexer's flattened decode path. Host metadata avoids a CUDA graph
403 + # sync when choosing the uniform or padded layout.
404 + per_req_lens = decode_metadata.per_req_decode_lens
405 + if per_req_lens is not None:
406 + use_uniform = (
407 + decode_metadata.decode_is_uniform
408 + and num_decode_tokens
409 + == num_requests * decode_metadata.write_max_decode_len
410 + )
411 + group_lens = per_req_lens
412 + lmax = decode_metadata.write_max_decode_len
413 + else:
414 + # Legacy metadata without per-request lens: fall back to the
415 + # host-side requires_padding flag. Unreached now (per-request
416 + # lens is always populated for decode), kept defensive.
417 + use_uniform = not decode_metadata.requires_padding
418 + group_lens = decode_metadata.decode_lens
419 + lmax = int(decode_metadata.decode_lens.max().item())
420 + if not use_uniform:
421 + # Non-uniform decode_lens (mixed plain-decode + spec-verify, or
422 + # a variable MTP-verify batch): scatter actual tokens into a
423 + # padded [B, lmax] layout. int32 tensors can't go through
424 + # pack_seq_triton (float/uint8 only). The scatter indices are
425 + # shared by all five scatters below (and the tail slot one).
426 + scatter_idx = _build_decode_scatter_indices(
427 + group_lens, num_requests, num_decode_tokens
428 + )
429 + dec_k = _scatter_decode_tokens_by_request(
430 + k[:num_decode_tokens], 0, num_requests, lmax, scatter_idx
431 + )
432 + dec_gate = _scatter_decode_tokens_by_request(
433 + gate_score[:num_decode_tokens],
434 + 0,
435 + num_requests,
436 + lmax,
437 + scatter_idx,
438 + )
439 + dec_slot = _scatter_decode_tokens_by_request(
440 + slot_mapping[:num_decode_tokens],
441 + -1,
442 + num_requests,
443 + lmax,
444 + scatter_idx,
445 + )
446 + dec_pos = _scatter_decode_tokens_by_request(
447 + positions[:num_decode_tokens].to(torch.int32),
448 + -1,
449 + num_requests,
450 + lmax,
451 + scatter_idx,
452 + )
453 + else:
454 + next_n = num_decode_tokens // num_requests
455 + shape2 = (num_requests, next_n)
456 + dec_k = k[:num_decode_tokens].view(*shape2, head_dim)
457 + dec_gate = gate_score[:num_decode_tokens].view(*shape2, head_dim)
458 + dec_slot = slot_mapping[:num_decode_tokens].view(shape2)
459 + dec_pos = positions[:num_decode_tokens].to(torch.int32).view(shape2)
460 + tail_meta = (
461 + attn_metadata.get(_resolve_layer_name(tail_prefix))
462 + if tail_prefix is not None
463 + else None
464 + )
465 + # Paged tail cache replaces the transient _DECODE_TAIL ring. Group
466 + # the tail group's token-granular slot_mapping per-request, mirroring
467 + # dec_slot / dec_pos, so the kernel gets each request's current-token
468 + # tail slot (block * kpool + pos % kpool).
469 + if tail_meta is not None:
470 + assert isinstance(tail_meta, DeepseekV32IndexerMetadata)
471 + if tail_meta is None or tail_kv_cache is None:
472 + dec_tail_slot = None
473 + elif not use_uniform:
474 + dec_tail_slot = _scatter_decode_tokens_by_request(
475 + tail_meta.slot_mapping[:num_decode_tokens],
476 + -1,
477 + num_requests,
478 + lmax,
479 + scatter_idx,
480 + )
481 + else:
482 + dec_tail_slot = tail_meta.slot_mapping[:num_decode_tokens].view(shape2)
483 + # The compress kernel writes the raw fp8 cache (not the quant view);
484 + # pass the underlying kv_cache, not kv_cache_quant_view.
485 + if dec_tail_slot is not None:
486 + # Single batched launch over [num_requests, next_n] replaces the
487 + # per-token sequential loop. The kernel iterates each request's
488 + # tokens in position order internally, preserving the
489 + # pool-completion read-after-stash dependency that the loop
490 + # provided. Inputs are already grouped per request (uniform:
491 + # view; non-uniform: _scatter_decode_tokens_by_request padded to
492 + # [B, lmax]) — no per-token .contiguous() copies needed.
493 + kpool_ops.kpool_decode_update_and_maybe_write_cache_batched(
494 + kv_cache_raw,
495 + tail_kv_cache,
496 + dec_tail_slot,
497 + dec_k,
498 + dec_gate,
499 + compress_ape,
500 + dec_slot,
501 + dec_pos,
502 + index_kpool,
503 + head_dim,
504 + round_scale=(scale_fmt is not None),
505 + )
506 + if current_platform.is_cuda_alike() and _fill_short_decode_causal_indices(
507 + topk_indices_buffer,
508 + positions,
509 + num_decode_tokens,
510 + attn_metadata_narrowed.max_seq_len,
511 + topk_tokens,
512 + ):
513 + return topk_indices_buffer
514 + decode_lens = decode_metadata.decode_lens
515 + if decode_metadata.requires_padding:
516 + # Padding also covers short chunked prefills classified as decode.
517 + # MXFP4 uses zero-byte padding so padded slots dequantize to zero.
518 + if q_scale is not None:
519 + padded_q_quant_decode_tokens = pack_seq_triton(
520 + q_quant[:num_decode_tokens], decode_lens, pad_value=0
521 + )
522 + padded_q_scale = pack_seq_triton(
523 + q_scale[:num_decode_tokens], decode_lens, pad_value=0
524 + )
525 + else:
526 + padded_q_quant_decode_tokens = pack_seq_triton(
527 + q_quant[:num_decode_tokens], decode_lens
528 + )
529 + padded_q_scale = None
530 + padded_weights = pack_seq_triton(
531 + weights[:num_decode_tokens], decode_lens, pad_value=0
532 + ).reshape(-1, *weights.shape[1:])
533 + else:
534 + padded_q_quant_decode_tokens = q_quant[:num_decode_tokens].reshape(
535 + decode_lens.shape[0], -1, *q_quant.shape[1:]
536 + )
537 + if q_scale is not None:
538 + padded_q_scale = q_scale[:num_decode_tokens].reshape(
539 + decode_lens.shape[0], -1, *q_scale.shape[1:]
540 + )
541 + else:
542 + padded_q_scale = None
543 + padded_weights = weights[:num_decode_tokens]
544 + # TODO: move and optimize below logic with triton kernels
545 + batch_size = padded_q_quant_decode_tokens.shape[0]
546 + next_n = padded_q_quant_decode_tokens.shape[1]
547 + num_padded_tokens = batch_size * next_n
548 + seq_lens = decode_metadata.seq_lens[:batch_size]
549 + # seq_lens is always 2D: (B, next_n) for native spec decode, (B, 1)
550 + # otherwise. deep_gemm fp8_fp4_paged_mqa_logits requires 2D context_lens;
551 + # the downstream topk kernels accept both 1D and 2D.
552 + padded_q_quant_cast = (
553 + padded_q_quant_decode_tokens.view(torch.int8)
554 + if use_fp4_cache
555 + else padded_q_quant_decode_tokens
556 + )
557 + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
558 + rocm_fp8_paged_mqa_logits,
559 + )
560 +
561 + assert padded_q_scale is None
562 + logits = rocm_fp8_paged_mqa_logits(
563 + padded_q_quant_cast,
564 + kv_cache,
565 + padded_weights[:num_padded_tokens],
566 + seq_lens,
567 + decode_metadata.block_table,
568 + decode_metadata.schedule_metadata,
569 + max_model_len=max_model_len,
570 + )
571 + num_rows = logits.shape[0]
572 + # kpool: logits are pool-granular -> select topk_tokens//kpool pools,
573 + # then expand each pool back to its kpool tokens.
574 + select_k = topk_tokens // index_kpool if index_kpool > 1 else topk_tokens
575 + if index_kpool > 1:
576 + pool_topk = torch.empty(
577 + (num_rows, select_k), dtype=torch.int32, device=logits.device
578 + )
579 + topk_dst = pool_topk
580 + else:
581 + topk_dst = topk_indices_buffer[:num_padded_tokens, :topk_tokens]
582 +
583 + torch.ops._C.top_k_per_row_decode(
584 + logits,
585 + next_n,
586 + seq_lens,
587 + topk_dst,
588 + num_rows,
589 + logits.stride(0),
590 + logits.stride(1),
591 + select_k,
592 + )
593 +
594 + # Resolve to token-level indices in the output buffer.
595 + if index_kpool > 1:
596 + pool_ids = pool_topk.to(torch.int64)
597 + n = pool_topk.shape[0]
598 + # Decode seq_lens are pool-granular; recover token lengths from
599 + # positions using the padded [B, next_n] row layout when needed.
600 + if positions is not None:
601 + dec_seq = _decode_topk_seq_lens(
602 + positions,
603 + decode_lens,
604 + num_decode_tokens,
605 + batch_size,
606 + next_n,
607 + decode_metadata.requires_padding,
608 + )
609 + else:
610 + dec_seq = decode_metadata.seq_lens[:n]
611 + if dec_seq.ndim == 2:
612 + dec_seq = dec_seq[:, -1]
613 + dec_seq = dec_seq.to(torch.int32)
614 + out = kpool_ops.expand_pools_and_append_tail(pool_ids, dec_seq, index_kpool)
615 + else:
616 + out = topk_dst
617 +
618 + if decode_metadata.requires_padding:
619 + # Drop padded query rows introduced by the next_n padding above.
620 + out = unpack_seq_triton(
621 + out.reshape(batch_size, -1, out.shape[-1]), decode_lens
622 + )
623 + topk_indices_buffer[: out.shape[0], : out.shape[-1]] = out
624 +
625 + return topk_indices_buffer
626 +
627 +
628 +@CustomOp.register("sparse_attn_indexer_kpool")
629 +class SparseAttnIndexerKpool(CustomOp):
630 + """Sparse Attention Indexer Custom Op Layer. This layer is extracted as a
631 + separate custom op since it involves heavy custom kernels like `mqa_logits`,
632 + `paged_mqa_logits` and `top_k_per_row`, etc. Those kernels maybe requires
633 + specific memory layout or implementation for different hardware backends to
634 + achieve optimal performance.
635 +
636 + For now, the default native path will use CUDA backend path. Other platform
637 + may requires add the corresponding Custom Op name `sparse_attn_indexer` to
638 + `custom_ops` in `CompilationConfig` to enable the platform specific path.
639 + """
640 +
641 + def __init__(
642 + self,
643 + k_cache,
644 + quant_block_size: int,
645 + scale_fmt: str,
646 + topk_tokens: int,
647 + head_dim: int,
648 + max_model_len: int,
649 + max_total_seq_len: int,
650 + topk_indices_buffer: torch.Tensor,
651 + skip_k_cache_insert: bool = False,
652 + use_fp4_cache: bool = False,
653 + tail_cache=None,
654 + ):
655 + super().__init__()
656 + self.k_cache = k_cache
657 + self.tail_cache = tail_cache
658 + self.quant_block_size = quant_block_size
659 + self.scale_fmt = scale_fmt
660 + self.topk_tokens = topk_tokens
661 + self.head_dim = head_dim
662 + self.max_model_len = max_model_len
663 + self.max_total_seq_len = max_total_seq_len
664 + self.topk_indices_buffer = topk_indices_buffer
665 + self.skip_k_cache_insert = skip_k_cache_insert
666 + self.use_fp4_cache = use_fp4_cache
667 +
668 + def forward_native(
669 + self,
670 + hidden_states: torch.Tensor,
671 + q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
672 + k: torch.Tensor,
673 + weights: torch.Tensor,
674 + *,
675 + gate_score: torch.Tensor | None = None,
676 + compress_ape: torch.Tensor | None = None,
677 + index_kpool: int = 1,
678 + positions: torch.Tensor | None = None,
679 + ):
680 + assert not self.use_fp4_cache, "AMD platform doesn't support fp4 cache yet"
681 + assert isinstance(q_quant, torch.Tensor), (
682 + "AMD sparse_attn_indexer expects a single FP8 q_quant tensor"
683 + )
684 + if not rocm_aiter_ops.is_enabled():
685 + raise RuntimeError(
686 + "Sparse attention indexer ROCm path is only supported on AITER. "
687 + "Please enable aiter with VLLM_ROCM_USE_AITER=1"
688 + )
689 + if index_kpool <= 1:
690 + return torch.ops.vllm.rocm_aiter_sparse_attn_indexer(
691 + hidden_states,
692 + _encode_layer_name(self.k_cache.prefix),
693 + self.k_cache.kv_cache,
694 + q_quant,
695 + k,
696 + weights,
697 + self.quant_block_size,
698 + self.scale_fmt,
699 + self.topk_tokens,
700 + self.head_dim,
701 + self.max_model_len,
702 + self.max_total_seq_len,
703 + self.topk_indices_buffer,
704 + skip_k_cache_insert=self.skip_k_cache_insert,
705 + )
706 + return sparse_attn_indexer_kpool(
707 + hidden_states,
708 + self.k_cache.prefix,
709 + self.k_cache.kv_cache,
710 + q_quant,
711 + None,
712 + k,
713 + weights,
714 + self.quant_block_size,
715 + self.scale_fmt,
716 + self.topk_tokens,
717 + self.head_dim,
718 + self.max_model_len,
719 + self.max_total_seq_len,
720 + self.topk_indices_buffer,
721 + self.skip_k_cache_insert,
722 + self.use_fp4_cache,
723 + gate_score,
724 + compress_ape,
725 + index_kpool,
726 + positions,
727 + self.tail_cache.kv_cache if self.tail_cache is not None else None,
728 + self.tail_cache.prefix if self.tail_cache is not None else None,
729 + )
vllm/models/glm5next/common/attention.py
+1 -1
@@ -23,7 +23,6 @@ from vllm.model_executor.layers.linear import (
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
26 from vllm.model_executor.models.deepseek_v2 import (
27 DeepSeekV2FusedQkvAProjLinear,
28 DeepseekV32IndexerCache,
@@ -31,6 +30,7 @@ from vllm.model_executor.models.deepseek_v2 import (
30 )
31 from vllm.model_executor.utils import maybe_disable_graph_partition
32 from vllm.models.glm5next.nvidia.ops.kpool_compress import fwht128_quant_fp8
33 +from vllm.models.glm5next.sparse_indexer import SparseAttnIndexerKpool
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
vllm/models/glm5next/common/sparse_indexer.py new
+170
@@ -0,0 +1,170 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Shared helpers for the glm5next sparse attention indexer (kpool) layers."""
4 +
5 +import torch
6 +
7 +RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024
8 +
9 +# MXFP4 layout: 2 values packed per byte, ue8m0 (1-byte) scale per block of 32.
10 +MXFP4_BLOCK_SIZE = 32
11 +
12 +
13 +def _build_decode_scatter_indices(
14 + decode_lens: torch.Tensor,
15 + num_requests: int,
16 + n: int,
17 +) -> tuple[torch.Tensor, torch.Tensor]:
18 + """Per-token (request id, intra-request index) for a non-uniform decode
19 + batch, with ``n == decode_lens.sum()`` as a host int (avoids a
20 + device sync and keeps both repeat_interleaves sync-free).
21 +
22 + Shared by every ``_scatter_decode_tokens_by_request`` call in a step:
23 + building it per call would repeat the same repeat_interleave/cumsum chain
24 + up to 5x per layer on the eager decode break.
25 + """
26 + device = decode_lens.device
27 + dl = decode_lens.to(torch.int64)
28 + req_id = torch.repeat_interleave(
29 + torch.arange(num_requests, device=device, dtype=torch.int64),
30 + dl,
31 + output_size=n,
32 + )
33 + req_starts = torch.cumsum(
34 + torch.cat([torch.zeros(1, device=device, dtype=torch.int64), dl[:-1]]),
35 + dim=0,
36 + )
37 + # Broadcast the per-request start offsets to per-token (length n ==
38 + # dl.sum()) so each token's intra-request index subtracts its own
39 + # request's start.
40 + starts = torch.repeat_interleave(req_starts, dl, output_size=n)
41 + intra = torch.arange(n, device=device, dtype=torch.int64) - starts
42 + return req_id, intra
43 +
44 +
45 +def _scatter_decode_tokens_by_request(
46 + tokens: torch.Tensor,
47 + pad_value,
48 + num_requests: int,
49 + lmax: int,
50 + scatter_indices: tuple[torch.Tensor, torch.Tensor],
51 +) -> torch.Tensor:
52 + """Group ``[N, ...]`` decode tokens into a padded ``[num_requests, lmax, ...]``
53 + layout: request ``r``'s tokens at row ``r`` in order; short requests padded.
54 +
55 + Unlike ``pack_seq_triton`` this is dtype-agnostic (needed for the int32
56 + slot/pos tensors) — it scatters with the shared per-step indices from
57 + ``_build_decode_scatter_indices``. Used only for the non-uniform
58 + (``requires_padding``) decode batch; uniform batches use a zero-copy
59 + reshape.
60 + """
61 + req_id, intra = scatter_indices
62 + out = torch.full(
63 + (num_requests, lmax, *tokens.shape[1:]),
64 + pad_value,
65 + dtype=tokens.dtype,
66 + device=tokens.device,
67 + )
68 + out[req_id, intra] = tokens
69 + return out
70 +
71 +
72 +def _decode_topk_seq_lens(
73 + positions: torch.Tensor,
74 + decode_lens: torch.Tensor,
75 + num_decode_tokens: int,
76 + batch_size: int,
77 + next_n: int,
78 + requires_padding: bool,
79 +) -> torch.Tensor:
80 + """Token-granular seq_len (pos + 1) per pool-topk row, layout-aware.
81 +
82 + ``pool_topk`` (and the logits it comes from) follow the padded
83 + ``[batch_size, next_n]`` grid whenever ``requires_padding`` is set, so row
84 + ``(b, t)`` corresponds to flat decode token ``offset_b + t`` -- NOT
85 + ``b * next_n + t``. Slicing flat ``positions[: batch_size * next_n]``
86 + (the uniform-layout shortcut) misaligns every row after the first
87 + non-uniform request and, past the decode region, reads prefill tokens'
88 + positions; ``expand_pools_and_append_tail`` then anchors the tail at
89 + another request's length, dropping the row's real tail tokens or emitting
90 + indices past its sequence (out-of-bounds block-table reads). Padded rows
91 + get 0 (empty tail); they are dropped by ``unpack_seq_triton`` anyway.
92 + """
93 + n = batch_size * next_n
94 + if not requires_padding:
95 + return positions[:n].to(torch.int32) + 1
96 + scatter_idx = _build_decode_scatter_indices(
97 + decode_lens, batch_size, num_decode_tokens
98 + )
99 + padded = _scatter_decode_tokens_by_request(
100 + positions[:num_decode_tokens].to(torch.int32),
101 + -1,
102 + batch_size,
103 + next_n,
104 + scatter_idx,
105 + )
106 + return padded.reshape(n) + 1 # pad rows: -1 + 1 = 0 -> empty tail
107 +
108 +
109 +def _fill_causal_indices(rows: torch.Tensor, positions: torch.Tensor) -> None:
110 + causal_range = torch.arange(rows.shape[1], device=rows.device, dtype=torch.int32)
111 + positions = positions.to(torch.int32)
112 + rows[:] = causal_range[None, :]
113 + rows[causal_range[None, :] > positions[:, None]] = -1
114 +
115 +
116 +def _fill_short_decode_causal_indices(
117 + topk_indices_buffer: torch.Tensor,
118 + positions: torch.Tensor | None,
119 + num_decode_tokens: int,
120 + max_seq_len: int,
121 + topk_tokens: int,
122 +) -> bool:
123 + """Fill exact causal rows when sparse decode would select every token."""
124 + if positions is None or positions.numel() == 0 or max_seq_len > topk_tokens:
125 + return False
126 + _fill_causal_indices(
127 + topk_indices_buffer[:num_decode_tokens], positions[:num_decode_tokens]
128 + )
129 + return True
130 +
131 +
132 +def _gather_workspace_shapes(
133 + total_seq_lens: int,
134 + head_dim: int,
135 + fp8_dtype: torch.dtype,
136 + use_fp4_cache: bool,
137 +) -> tuple[tuple[tuple[int, int], torch.dtype], tuple[tuple[int, int], torch.dtype]]:
138 + """Return ((values_shape, values_dtype), (scales_shape, scales_dtype)) for
139 + the K-gather workspace. FP8 path: (T, head_dim) fp8 + (T, 4) uint8 fp32
140 + scales. MXFP4 path: (T, head_dim // 2) uint8 packed mxfp4 +
141 + (T, head_dim // MXFP4_BLOCK_SIZE) uint8 ue8m0 scales."""
142 + if use_fp4_cache:
143 + return (
144 + ((total_seq_lens, head_dim // 2), torch.uint8),
145 + ((total_seq_lens, head_dim // MXFP4_BLOCK_SIZE), torch.uint8),
146 + )
147 + return (
148 + ((total_seq_lens, head_dim), fp8_dtype),
149 + ((total_seq_lens, 4), torch.uint8),
150 + )
151 +
152 +
153 +def kv_cache_as_quant_view(
154 + kv_cache: torch.Tensor,
155 + head_dim: int,
156 + use_fp4_cache: bool,
157 +) -> torch.Tensor:
158 + """4D ``[num_blocks, block_size, 1, head_width]`` view expected by
159 + DeepGEMM, from the 3D indexer kv-cache allocation."""
160 + if use_fp4_cache:
161 + assert kv_cache.ndim == 3 and kv_cache.dtype == torch.uint8
162 + num_blocks, block_size, _ = kv_cache.shape
163 + page_bytes = int(kv_cache.stride(0))
164 + fp4_bytes = head_dim // 2 + head_dim // MXFP4_BLOCK_SIZE
165 + return torch.as_strided(
166 + kv_cache,
167 + size=(num_blocks, block_size, 1, fp4_bytes),
168 + stride=(page_bytes, fp4_bytes, fp4_bytes, 1),
169 + )
170 + return kv_cache.unsqueeze(-2)
vllm/models/glm5next/nvidia/sparse_indexer.py renamed
+76 -391
@@ -2,30 +2,29 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Custom Sparse Attention Indexer layers."""
4
5 -from typing import TYPE_CHECKING
6 -
5 import torch
6
7 import vllm.envs as envs
10 -from vllm._aiter_ops import rocm_aiter_ops
8 from vllm.compilation.breakable_cudagraph import eager_break_during_capture
9 from vllm.config import get_current_vllm_config_or_none
10 from vllm.forward_context import get_forward_context
11 from vllm.logger import init_logger
12 from vllm.model_executor.custom_op import CustomOp
13 +from vllm.models.glm5next.common.sparse_indexer import (
14 + RADIX_TOPK_WORKSPACE_SIZE,
15 + _build_decode_scatter_indices,
16 + _decode_topk_seq_lens,
17 + _fill_causal_indices,
18 + _fill_short_decode_causal_indices,
19 + _gather_workspace_shapes,
20 + _scatter_decode_tokens_by_request,
21 + kv_cache_as_quant_view,
22 +)
23 +from vllm.models.glm5next.nvidia.ops import kpool_compress as kpool_ops
24 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,
28 _resolve_layer_name,
29 )
30 from vllm.v1.attention.backends.mla.indexer import (
@@ -36,16 +35,9 @@ from vllm.v1.worker.workspace import current_workspace_manager
35
36 if current_platform.is_cuda_alike():
37 from vllm import _custom_ops as ops
39 -elif current_platform.is_xpu():
40 - from vllm._xpu_ops import xpu_ops
38
39 logger = init_logger(__name__)
40
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 -
41 # kpool write helper: form pools from the current token batch and compress them
42 # into the index K cache via the fused Triton kernel.
43
@@ -96,166 +88,6 @@ def _kpool_compress_insert(
88 )
89
90
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 -
91 @eager_break_during_capture
92 def sparse_attn_indexer_kpool(
93 hidden_states: torch.Tensor,
@@ -393,26 +225,13 @@ def sparse_attn_indexer_kpool(
225 else:
226 # standard: per-token fp8 quant + scatter (all tokens).
227 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 - )
228 + ops.indexer_k_quant_and_cache(
229 + k,
230 + kv_cache,
231 + slot_mapping,
232 + quant_block_size,
233 + scale_fmt,
234 + )
235
236 topk_indices_buffer[: hidden_states.shape[0]] = -1
237 if has_prefill:
@@ -466,27 +285,13 @@ def sparse_attn_indexer_kpool(
285 k_scale = k_scale_full[: chunk.total_seq_lens]
286
287 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 - )
288 + ops.cp_gather_indexer_k_quant_cache(
289 + kv_cache,
290 + k_quant,
291 + k_scale,
292 + chunk.block_table,
293 + chunk.cu_seq_lens,
294 + )
295
296 q_slice = q_quant[chunk.token_start : chunk.token_end]
297 q_scale_slice = (
@@ -504,30 +309,16 @@ def sparse_attn_indexer_kpool(
309 q_slice_cast = q_slice
310 k_quant_cast = k_quant
311 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 - )
312 + from vllm.utils.deep_gemm import fp8_fp4_mqa_logits
313 +
314 + logits = fp8_fp4_mqa_logits(
315 + (q_slice_cast, q_scale_slice),
316 + (k_quant_cast, k_scale_cast),
317 + weights[chunk.token_start : chunk.token_end],
318 + chunk.cu_seqlen_ks,
319 + chunk.cu_seqlen_ke,
320 + clean_logits=False,
321 + )
322 num_rows = logits.shape[0]
323
324 # kpool: logits are pool-granular (compress_ratio == index_kpool),
@@ -544,28 +335,16 @@ def sparse_attn_indexer_kpool(
335 chunk.token_start : chunk.token_end, :topk_tokens
336 ]
337
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 - )
338 + torch.ops._C.top_k_per_row_prefill(
339 + logits,
340 + chunk.cu_seqlen_ks,
341 + chunk.cu_seqlen_ke,
342 + topk_dst,
343 + num_rows,
344 + logits.stride(0),
345 + logits.stride(1),
346 + select_k,
347 + )
348
349 if index_kpool > 1:
350 pool_ids = pool_topk.to(torch.int64)
@@ -765,34 +544,18 @@ def sparse_attn_indexer_kpool(
544 if use_fp4_cache
545 else padded_q_quant_decode_tokens
546 )
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
547 + from vllm.utils.deep_gemm import fp8_fp4_paged_mqa_logits
548
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 - )
549 + logits = fp8_fp4_paged_mqa_logits(
550 + (padded_q_quant_cast, padded_q_scale),
551 + kv_cache,
552 + padded_weights[:num_padded_tokens],
553 + seq_lens,
554 + decode_metadata.block_table,
555 + decode_metadata.schedule_metadata,
556 + max_model_len=max_model_len,
557 + clean_logits=False,
558 + )
559 num_rows = logits.shape[0]
560 # kpool: logits are pool-granular -> select topk_tokens//kpool pools,
561 # then expand each pool back to its kpool tokens.
@@ -819,28 +582,16 @@ def sparse_attn_indexer_kpool(
582 attn_metadata_narrowed.max_seq_len,
583 )
584 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 - )
585 + torch.ops._C.top_k_per_row_decode(
586 + logits,
587 + next_n,
588 + seq_lens,
589 + topk_dst,
590 + num_rows,
591 + logits.stride(0),
592 + logits.stride(1),
593 + select_k,
594 + )
595
596 # Resolve to token-level indices in the output buffer.
597 if index_kpool > 1:
@@ -942,33 +693,16 @@ class SparseAttnIndexerKpool(CustomOp):
693 index_kpool: int = 1,
694 positions: torch.Tensor | None = None,
695 ):
945 - if current_platform.is_cuda() or current_platform.is_xpu():
946 - return self.forward_cuda(
947 - hidden_states,
948 - q_quant,
949 - k,
950 - weights,
951 - gate_score=gate_score,
952 - compress_ape=compress_ape,
953 - index_kpool=index_kpool,
954 - positions=positions,
955 - )
956 - elif current_platform.is_rocm():
957 - return self.forward_hip(
958 - hidden_states,
959 - q_quant,
960 - k,
961 - weights,
962 - gate_score=gate_score,
963 - compress_ape=compress_ape,
964 - index_kpool=index_kpool,
965 - positions=positions,
966 - )
967 - else:
968 - raise NotImplementedError(
969 - "SparseAttnIndexer native forward is only implemented for "
970 - "CUDA, ROCm and XPU platforms."
971 - )
696 + return self.forward_cuda(
697 + hidden_states,
698 + q_quant,
699 + k,
700 + weights,
701 + gate_score=gate_score,
702 + compress_ape=compress_ape,
703 + index_kpool=index_kpool,
704 + positions=positions,
705 + )
706
707 def forward_cuda(
708 self,
@@ -1012,52 +746,3 @@ class SparseAttnIndexerKpool(CustomOp):
746 self.tail_cache.kv_cache if self.tail_cache is not None else None,
747 self.tail_cache.prefix if self.tail_cache is not None else None,
748 )
1015 -
1016 - def forward_hip(
1017 - self,
1018 - hidden_states: torch.Tensor,
1019 - q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
1020 - k: torch.Tensor,
1021 - weights: torch.Tensor,
1022 - *,
1023 - gate_score: torch.Tensor | None = None,
1024 - compress_ape: torch.Tensor | None = None,
1025 - index_kpool: int = 1,
1026 - positions: torch.Tensor | None = None,
1027 - ):
1028 - assert not self.use_fp4_cache, "AMD platform doesn't support fp4 cache yet"
1029 - assert isinstance(q_quant, torch.Tensor), (
1030 - "AMD sparse_attn_indexer expects a single FP8 q_quant tensor"
1031 - )
1032 - if rocm_aiter_ops.is_enabled():
1033 - if index_kpool <= 1:
1034 - return torch.ops.vllm.rocm_aiter_sparse_attn_indexer(
1035 - hidden_states,
1036 - _encode_layer_name(self.k_cache.prefix),
1037 - self.k_cache.kv_cache,
1038 - q_quant,
1039 - k,
1040 - weights,
1041 - self.quant_block_size,
1042 - self.scale_fmt,
1043 - self.topk_tokens,
1044 - self.head_dim,
1045 - self.max_model_len,
1046 - self.max_total_seq_len,
1047 - self.topk_indices_buffer,
1048 - skip_k_cache_insert=self.skip_k_cache_insert,
1049 - )
1050 - return self.forward_cuda(
1051 - hidden_states,
1052 - q_quant,
1053 - k,
1054 - weights,
1055 - gate_score=gate_score,
1056 - compress_ape=compress_ape,
1057 - index_kpool=index_kpool,
1058 - positions=positions,
1059 - )
1060 - raise RuntimeError(
1061 - "Sparse attention indexer ROCm path is only supported on AITER. "
1062 - "Please enable aiter with VLLM_ROCM_USE_AITER=1"
1063 - )
vllm/models/glm5next/sparse_indexer.py new
+22
@@ -0,0 +1,22 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Sparse attention indexer (kpool) — hardware-isolated entry point.
4 +
5 +The implementation lives under ``nvidia/`` and ``amd/`` (with shared helpers
6 +in ``common/``); this module picks the right one for the current platform.
7 +"""
8 +
9 +from vllm.platforms import current_platform
10 +
11 +if current_platform.is_rocm():
12 + from .amd.sparse_indexer import (
13 + SparseAttnIndexerKpool,
14 + )
15 +else:
16 + from .nvidia.sparse_indexer import (
17 + SparseAttnIndexerKpool,
18 + )
19 +
20 +__all__ = [
21 + "SparseAttnIndexerKpool",
22 +]