| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """GLM-5.3-Flash KDA recurrent (decode) kernel. |
| 4 | |
| 5 | The decode path hands the kernel column slices of the merged ``q|k|v`` conv |
| 6 | output and of the fused ``qkvbfg_a`` projection (beta), so q/k/v/beta are |
| 7 | token-strided rather than contiguous. Both backends must match a pure-PyTorch |
| 8 | recurrence, read these slices in place and reject unsupported layouts. |
| 9 | """ |
| 10 | |
| 11 | import pytest |
| 12 | import torch |
| 13 | |
| 14 | from vllm.platforms import current_platform |
| 15 | |
| 16 | if current_platform.is_rocm(): |
| 17 | from vllm.models.glm5next.amd.ops.third_party.kda import fused_recurrent_kda |
| 18 | else: |
| 19 | from vllm.models.glm5next.nvidia.ops.third_party.kda import fused_recurrent_kda |
| 20 | |
| 21 | pytestmark = pytest.mark.skipif( |
| 22 | not current_platform.is_cuda_alike(), reason="Requires CUDA or ROCm" |
| 23 | ) |
| 24 | |
| 25 | H, D = 16, 128 |
| 26 | LOWER_BOUND = -5.0 |
| 27 | |
| 28 | |
| 29 | def naive_recurrent_kda( |
| 30 | q: torch.Tensor, |
| 31 | k: torch.Tensor, |
| 32 | v: torch.Tensor, |
| 33 | raw_g: torch.Tensor, |
| 34 | raw_beta: torch.Tensor, |
| 35 | a_log: torch.Tensor, |
| 36 | g_bias: torch.Tensor, |
| 37 | state: torch.Tensor, |
| 38 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 39 | """fp32 reference for one sequence: ``[T, H, D]`` inputs, ``[H, D, D]`` |
| 40 | (v-major) state; mirrors the kernel's in-kernel gate, beta sigmoid and |
| 41 | q/k l2norm. |
| 42 | """ |
| 43 | q, k, v, raw_g, raw_beta = (x.float() for x in (q, k, v, raw_g, raw_beta)) |
| 44 | q = q / torch.sqrt(q.square().sum(-1, keepdim=True) + 1e-6) * D**-0.5 |
| 45 | k = k / torch.sqrt(k.square().sum(-1, keepdim=True) + 1e-6) |
| 46 | gate = LOWER_BOUND * torch.sigmoid(a_log.exp()[:, None] * (raw_g + g_bias)) |
| 47 | beta = torch.sigmoid(raw_beta) |
| 48 | s = state.clone() |
| 49 | out = torch.empty_like(v) |
| 50 | for t in range(q.shape[0]): |
| 51 | s = s * gate[t].exp()[:, None, :] |
| 52 | u = beta[t][:, None] * (v[t] - torch.einsum("hvk,hk->hv", s, k[t])) |
| 53 | s = s + u[:, :, None] * k[t][:, None, :] |
| 54 | out[t] = torch.einsum("hvk,hk->hv", s, q[t]) |
| 55 | return out, s |
| 56 | |
| 57 | |
| 58 | def make_inputs(num_seqs: int, query_len: int, device: torch.device): |
| 59 | """Token-strided q/k/v/beta as the decode path produces them: column |
| 60 | slices of a merged ``[T, q|k|v]`` conv output and of the fused |
| 61 | ``[T, qkv|beta|f_a|g_a]`` projection. |
| 62 | """ |
| 63 | T, proj = num_seqs * query_len, H * D |
| 64 | qkv = torch.randn(T, 3 * proj, dtype=torch.bfloat16, device=device) |
| 65 | projected = torch.randn( |
| 66 | T, 3 * proj + H + 2 * D, dtype=torch.bfloat16, device=device |
| 67 | ) |
| 68 | q, k, v = (x.view(1, T, H, D) for x in qkv.split(proj, dim=-1)) |
| 69 | beta = projected[:, 3 * proj : 3 * proj + H].unsqueeze(0) |
| 70 | # (A size-1 token dim gets an arbitrary stride from `view`.) |
| 71 | assert T == 1 or (q.stride(1) == 3 * proj and beta.stride(1) == projected.stride(0)) |
| 72 | inputs = dict( |
| 73 | q=q, |
| 74 | k=k, |
| 75 | v=v, |
| 76 | g=torch.randn(1, T, H, D, dtype=torch.bfloat16, device=device), |
| 77 | beta=beta, |
| 78 | a_log=0.5 * torch.randn(H, dtype=torch.float32, device=device), |
| 79 | g_bias=0.1 * torch.randn(H * D, dtype=torch.float32, device=device), |
| 80 | cu_seqlens=torch.arange(0, T + 1, query_len, dtype=torch.int32, device=device), |
| 81 | ) |
| 82 | # Slot 0 is NULL_BLOCK_ID; sequences own random distinct slots (one per |
| 83 | # token in the spec-decode layout). |
| 84 | slots = torch.randperm(T, device=device).to(torch.int32) + 1 |
| 85 | if query_len == 1: |
| 86 | inputs["ssm_state_indices"] = slots |
| 87 | else: |
| 88 | inputs["ssm_state_indices"] = slots.view(num_seqs, query_len) |
| 89 | inputs["num_accepted_tokens"] = torch.randint( |
| 90 | 1, query_len + 1, (num_seqs,), dtype=torch.int32, device=device |
| 91 | ) |
| 92 | state = torch.randn(T + 1, H, D, D, dtype=torch.float32, device=device) |
| 93 | return inputs, state |
| 94 | |
| 95 | |
| 96 | def run_kernel(inputs: dict, state: torch.Tensor) -> torch.Tensor: |
| 97 | out, _ = fused_recurrent_kda( |
| 98 | **inputs, |
| 99 | initial_state=state, |
| 100 | use_qk_l2norm_in_kernel=True, |
| 101 | sigmoid_beta=True, |
| 102 | compute_gate=True, |
| 103 | lower_bound=LOWER_BOUND, |
| 104 | ) |
| 105 | return out |
| 106 | |
| 107 | |
| 108 | @pytest.mark.parametrize( |
| 109 | ("num_seqs", "query_len"), [(1, 1), (7, 1), (3, 3)], ids=["1x1", "7x1", "3x3"] |
| 110 | ) |
| 111 | @torch.inference_mode() |
| 112 | def test_fused_recurrent_kda_matches_reference(num_seqs: int, query_len: int): |
| 113 | torch.manual_seed(0) |
| 114 | device = torch.device("cuda") |
| 115 | inputs, state = make_inputs(num_seqs, query_len, device) |
| 116 | expected_state = state.clone() |
| 117 | out = run_kernel(inputs, state) |
| 118 | |
| 119 | indices = inputs["ssm_state_indices"].view(num_seqs, query_len) |
| 120 | accepted = inputs.get("num_accepted_tokens") |
| 121 | expected = torch.empty_like(out[0], dtype=torch.float32) |
| 122 | for n in range(num_seqs): |
| 123 | first = indices[n, 0 if accepted is None else accepted[n] - 1] |
| 124 | s = expected_state[first] |
| 125 | for t in range(query_len): |
| 126 | tok = slice(n * query_len + t, n * query_len + t + 1) |
| 127 | expected[tok], s = naive_recurrent_kda( |
| 128 | inputs["q"][0, tok], |
| 129 | inputs["k"][0, tok], |
| 130 | inputs["v"][0, tok], |
| 131 | inputs["g"][0, tok], |
| 132 | inputs["beta"][0, tok], |
| 133 | inputs["a_log"], |
| 134 | inputs["g_bias"].view(H, D), |
| 135 | s, |
| 136 | ) |
| 137 | expected_state[indices[n, t]] = s |
| 138 | |
| 139 | torch.testing.assert_close(out[0].float(), expected, rtol=1e-2, atol=1e-3) |
| 140 | used = indices.flatten().long() |
| 141 | torch.testing.assert_close(state[used], expected_state[used], rtol=1e-5, atol=1e-5) |
| 142 | |
| 143 | |
| 144 | @pytest.mark.parametrize( |
| 145 | ("num_seqs", "query_len"), [(7, 1), (3, 3)], ids=["7x1", "3x3"] |
| 146 | ) |
| 147 | @torch.inference_mode() |
| 148 | def test_fused_recurrent_kda_strided_inputs_bit_identical_to_contiguous( |
| 149 | num_seqs: int, query_len: int |
| 150 | ): |
| 151 | torch.manual_seed(0) |
| 152 | device = torch.device("cuda") |
| 153 | inputs, state = make_inputs(num_seqs, query_len, device) |
| 154 | for name in ("q", "k", "v", "beta"): |
| 155 | assert not inputs[name].is_contiguous() |
| 156 | contiguous = { |
| 157 | name: x.contiguous() if name in ("q", "k", "v", "beta") else x |
| 158 | for name, x in inputs.items() |
| 159 | } |
| 160 | state_ref = state.clone() |
| 161 | out_ref = run_kernel(contiguous, state_ref) |
| 162 | out = run_kernel(inputs, state) |
| 163 | torch.testing.assert_close(out, out_ref, rtol=0, atol=0) |
| 164 | torch.testing.assert_close(state, state_ref, rtol=0, atol=0) |
| 165 | |
| 166 | |
| 167 | @pytest.mark.skipif( |
| 168 | not current_platform.is_cuda_alike(), |
| 169 | reason="Both implementations require zero-copy token-strided layouts", |
| 170 | ) |
| 171 | @torch.inference_mode() |
| 172 | def test_fused_recurrent_kda_rejects_unaddressable_layouts(): |
| 173 | """Layouts the token-stride addressing cannot express must fail loudly |
| 174 | rather than read the wrong tokens: a batch slice of a wider buffer |
| 175 | (``stride(0) != T * stride(1)``), overlapping tokens, and a head-strided |
| 176 | (transposed) block. |
| 177 | """ |
| 178 | torch.manual_seed(0) |
| 179 | device = torch.device("cuda") |
| 180 | inputs, state = make_inputs(1, 1, device) |
| 181 | T = 4 |
| 182 | base = torch.randn(4, T, 2 * H * D, dtype=torch.bfloat16, device=device) |
| 183 | bad_q = { |
| 184 | "batch slice": base[::2, :, : H * D].view(2, T, H, D), |
| 185 | "overlapping tokens": base[:1, :, : H * D].as_strided( |
| 186 | (1, T, H, D), (0, D, D, 1) |
| 187 | ), |
| 188 | "head-strided": base[:1, :, : H * D].view(1, T, D, H).transpose(2, 3), |
| 189 | } |
| 190 | for q in bad_q.values(): |
| 191 | broken = dict(inputs, q=q, k=q, v=q) |
| 192 | broken["cu_seqlens"] = None if q.shape[0] > 1 else inputs["cu_seqlens"] |
| 193 | with pytest.raises(AssertionError, match=r"torch.Size"): |
| 194 | run_kernel(broken, state) |
| 195 | |
| 196 | |
| 197 | @pytest.mark.parametrize("offset", [16, 144, 256]) |
| 198 | @pytest.mark.parametrize("dim_first", [False, True]) |
| 199 | @pytest.mark.parametrize("num_spec", [0, 3]) |
| 200 | @torch.inference_mode() |
| 201 | def test_prefill_checkpoint_resumes_suffix(monkeypatch, dim_first, num_spec, offset): |
| 202 | """Restoring both cached states must reproduce an uninterrupted prefill.""" |
| 203 | from types import SimpleNamespace |
| 204 | |
| 205 | from vllm.model_executor.layers.mamba.checkpoint import ( |
| 206 | MambaPrefillCheckpointMetadata, |
| 207 | ) |
| 208 | from vllm.model_executor.layers.mamba.kda_checkpoint import ( |
| 209 | FlashKDAPrefillCheckpointExporter, |
| 210 | ) |
| 211 | from vllm.models.glm5next.common import kda |
| 212 | from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata |
| 213 | from vllm.v1.attention.backends.utils import compute_causal_conv1d_metadata |
| 214 | from vllm.v1.worker.workspace import WorkspaceManager |
| 215 | |
| 216 | if kda._resolve_kda_prefill_backend("auto", D, torch.bfloat16, -5) != "flashkda": |
| 217 | pytest.skip("requires FlashKDA") |
| 218 | import vllm._flashkda_C # noqa: F401 |
| 219 | |
| 220 | torch.manual_seed(42) |
| 221 | device = torch.device("cuda") |
| 222 | length, heads = 273, 2 |
| 223 | width, state_len = 3 * heads * D, 3 + num_spec |
| 224 | layer = object.__new__(kda.Glm5NextLinearAttention) |
| 225 | torch.nn.Module.__init__(layer) |
| 226 | layer.prefix = "layer.0" |
| 227 | layer.local_num_heads = heads |
| 228 | layer.head_dim = D |
| 229 | layer.conv_size = 4 |
| 230 | layer.local_projection_size = heads * D |
| 231 | layer.kda_safe_gate = True |
| 232 | layer.kda_lower_bound = -5.0 |
| 233 | layer.kda_prefill_backend = "flashkda" |
| 234 | layer._conv_state_dim_first = dim_first |
| 235 | layer._merged_conv_weight = torch.randn(width, 4, device=device) * 0.1 |
| 236 | layer.q_conv1d = SimpleNamespace(bias=None) |
| 237 | layer.A_log = torch.randn(1, 1, heads, 1, device=device) * 0.1 |
| 238 | layer.dt_bias = torch.randn(heads * D, device=device) * 0.1 |
| 239 | state_shape = (2, heads, D, D) |
| 240 | layer._flashkda_buffer_specs = ( |
| 241 | (state_shape, torch.float32), |
| 242 | (state_shape, torch.float32), |
| 243 | ( |
| 244 | (torch.ops._flashkda_C.get_workspace_size(2 * length, heads, 2),), |
| 245 | torch.uint8, |
| 246 | ), |
| 247 | ((1, 2 * length, heads, D), torch.bfloat16), |
| 248 | ) |
| 249 | layer._checkpoint_exporter = FlashKDAPrefillCheckpointExporter(state_len=3) |
| 250 | workspace = WorkspaceManager(device) |
| 251 | monkeypatch.setattr(kda, "current_workspace_manager", lambda: workspace) |
| 252 | conv_shape = (width, state_len) if dim_first else (state_len, width) |
| 253 | num_slots = 5 + (num_spec + 1 if num_spec else 0) |
| 254 | conv_storage = torch.zeros( |
| 255 | num_slots, width * state_len + 8, device=device, dtype=torch.bfloat16 |
| 256 | ) |
| 257 | conv = conv_storage[:, : width * state_len].view(num_slots, *conv_shape) |
| 258 | recurrent_storage = torch.zeros( |
| 259 | num_slots, heads * D * D + 8, device=device, dtype=torch.float32 |
| 260 | ) |
| 261 | recurrent = recurrent_storage[:, : heads * D * D].view(num_slots, heads, D, D) |
| 262 | layer.kv_cache = (conv, recurrent) |
| 263 | projected = torch.randn( |
| 264 | 2 * length, width + heads, device=device, dtype=torch.bfloat16 |
| 265 | ) |
| 266 | raw_qkv, beta = projected[:, :width], projected[:, width:].unsqueeze(0) |
| 267 | raw_before = raw_qkv.clone() |
| 268 | gate = torch.randn(1, 2 * length, heads, D, device=device, dtype=torch.bfloat16) |
| 269 | |
| 270 | def run(qkv, g, b, slots, has_initial, checkpoint=None): |
| 271 | n = qkv.shape[0] // 2 |
| 272 | cu_cpu = torch.tensor([0, n, 2 * n], dtype=torch.int32) |
| 273 | nums, batch_ptr, chunk_offsets = compute_causal_conv1d_metadata( |
| 274 | cu_cpu, device=device |
| 275 | ) |
| 276 | metadata = GDNAttentionMetadata( |
| 277 | num_prefills=2, |
| 278 | num_prefill_tokens=2 * n, |
| 279 | num_decodes=0, |
| 280 | num_decode_tokens=0, |
| 281 | num_spec_decodes=0, |
| 282 | num_spec_decode_tokens=0, |
| 283 | num_actual_tokens=2 * n, |
| 284 | has_initial_state=torch.tensor(has_initial, device=device), |
| 285 | non_spec_query_start_loc=cu_cpu.to(device), |
| 286 | non_spec_state_indices_tensor=torch.tensor( |
| 287 | slots, dtype=torch.int32, device=device |
| 288 | ), |
| 289 | checkpoint=checkpoint, |
| 290 | nums_dict=nums, |
| 291 | batch_ptr=batch_ptr, |
| 292 | token_chunk_offset_ptr=chunk_offsets, |
| 293 | ) |
| 294 | monkeypatch.setattr( |
| 295 | kda, |
| 296 | "get_forward_context", |
| 297 | lambda: SimpleNamespace(attn_metadata={layer.prefix: metadata}), |
| 298 | ) |
| 299 | spec_tokens = num_spec + 1 if num_spec else 0 |
| 300 | if spec_tokens: |
| 301 | qkv = torch.cat([qkv[:spec_tokens], qkv]) |
| 302 | g = torch.cat([g[:, :spec_tokens], g], dim=1) |
| 303 | b = torch.cat([b[:, :spec_tokens], b], dim=1) |
| 304 | metadata.num_spec_decodes = 1 |
| 305 | metadata.num_spec_decode_tokens = spec_tokens |
| 306 | metadata.num_actual_tokens += spec_tokens |
| 307 | metadata.spec_sequence_masks = torch.tensor( |
| 308 | [True, False, False], device=device |
| 309 | ) |
| 310 | metadata.spec_query_start_loc = torch.tensor( |
| 311 | [0, spec_tokens], dtype=torch.int32, device=device |
| 312 | ) |
| 313 | metadata.spec_state_indices_tensor = torch.arange( |
| 314 | 5, num_slots, dtype=torch.int32, device=device |
| 315 | ).unsqueeze(0) |
| 316 | metadata.spec_token_indx = torch.arange(spec_tokens, device=device) |
| 317 | metadata.non_spec_token_indx = torch.arange( |
| 318 | spec_tokens, qkv.shape[0], device=device |
| 319 | ) |
| 320 | metadata.num_accepted_tokens = torch.ones( |
| 321 | 1, dtype=torch.int32, device=device |
| 322 | ) |
| 323 | out = torch.empty( |
| 324 | 1, qkv.shape[0], heads, D, device=device, dtype=torch.bfloat16 |
| 325 | ) |
| 326 | layer._forward(qkv, g, b, out) |
| 327 | return out[:, spec_tokens:] |
| 328 | |
| 329 | checkpoint = MambaPrefillCheckpointMetadata( |
| 330 | torch.tensor([offset, 0], dtype=torch.int32, device=device), |
| 331 | torch.tensor([3, 0], dtype=torch.int32, device=device), |
| 332 | ) |
| 333 | expected = run(raw_qkv, gate, beta, [1, 2], [False, False], checkpoint) |
| 334 | expected_state = recurrent[1].clone() |
| 335 | saved_conv, saved_recurrent = conv[3].clone(), recurrent[3].clone() |
| 336 | torch.testing.assert_close(raw_qkv, raw_before) |
| 337 | conv_view = conv if dim_first else conv.transpose(-1, -2) |
| 338 | torch.testing.assert_close(conv_view[3, :, :3], raw_qkv[offset - 3 : offset].T) |
| 339 | assert torch.count_nonzero(conv_view[3, :, 3:]) == 0 |
| 340 | assert torch.count_nonzero(conv[0]) == torch.count_nonzero(recurrent[0]) == 0 |
| 341 | |
| 342 | run( |
| 343 | torch.cat([raw_qkv[:offset], raw_qkv[length : length + offset]]), |
| 344 | torch.cat([gate[:, :offset], gate[:, length : length + offset]], dim=1), |
| 345 | torch.cat([beta[:, :offset], beta[:, length : length + offset]], dim=1), |
| 346 | [1, 2], |
| 347 | [False, False], |
| 348 | ) |
| 349 | torch.testing.assert_close(saved_recurrent, recurrent[1], atol=0.003, rtol=0.03) |
| 350 | |
| 351 | suffix = slice(offset, length) |
| 352 | qkv = torch.cat([raw_qkv[suffix], raw_qkv[length + offset :]]) |
| 353 | g = torch.cat([gate[:, suffix], gate[:, length + offset :]], dim=1) |
| 354 | b = torch.cat([beta[:, suffix], beta[:, length + offset :]], dim=1) |
| 355 | conv[4].copy_(saved_conv) |
| 356 | recurrent[4].copy_(saved_recurrent) |
| 357 | actual = run(qkv, g, b, [3, 4], [True, True]) |
| 358 | torch.testing.assert_close( |
| 359 | actual[:, : length - offset], expected[:, suffix], atol=0.003, rtol=0.03 |
| 360 | ) |
| 361 | torch.testing.assert_close(recurrent[3], expected_state, atol=0.003, rtol=0.03) |