| 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) |