| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """The fused ROCm KDA chunk kernel must match the Triton chunk path. |
| 4 | |
| 5 | The fused kernel reassociates nothing, but it keeps the chunk states in |
| 6 | registers and rounds them to bfloat16 only where the Triton path does, so the |
| 7 | outputs agree to bfloat16 rounding and the final states nearly exactly. |
| 8 | """ |
| 9 | |
| 10 | import pytest |
| 11 | import torch |
| 12 | |
| 13 | from vllm.platforms import current_platform |
| 14 | |
| 15 | |
| 16 | def _fused_chunk_available() -> bool: |
| 17 | if not (current_platform.is_rocm() and torch.cuda.is_available()): |
| 18 | return False |
| 19 | from vllm.models.kimi_k3.amd.ops.kda_chunk import is_fused_kda_chunk_supported |
| 20 | |
| 21 | return is_fused_kda_chunk_supported() |
| 22 | |
| 23 | |
| 24 | pytestmark = pytest.mark.skipif( |
| 25 | not _fused_chunk_available(), |
| 26 | reason="The fused KDA chunk kernel needs a gfx950 build that includes it", |
| 27 | ) |
| 28 | |
| 29 | HEAD_DIM = 128 |
| 30 | NUM_HEADS = 12 |
| 31 | LOWER_BOUND = -5.0 |
| 32 | |
| 33 | |
| 34 | def _requires_kernel() -> None: |
| 35 | if not hasattr(torch.ops._C, "fused_kda_chunk"): |
| 36 | pytest.skip("vLLM was built without the fused KDA chunk kernel") |
| 37 | |
| 38 | |
| 39 | def _inputs(seqlens: list[int], seed: int) -> dict: |
| 40 | torch.manual_seed(seed) |
| 41 | dev, dtype = "cuda", torch.bfloat16 |
| 42 | total = sum(seqlens) |
| 43 | cu = torch.tensor( |
| 44 | [0] + torch.tensor(seqlens).cumsum(0).tolist(), device=dev, dtype=torch.int32 |
| 45 | ) |
| 46 | shape = (1, total, NUM_HEADS, HEAD_DIM) |
| 47 | return dict( |
| 48 | q=torch.randn(shape, device=dev, dtype=dtype) * 0.5, |
| 49 | k=torch.randn(shape, device=dev, dtype=dtype) * 0.5, |
| 50 | v=torch.randn(shape, device=dev, dtype=dtype) * 0.5, |
| 51 | raw_g=torch.randn(shape, device=dev, dtype=dtype), |
| 52 | raw_beta=torch.randn(1, total, NUM_HEADS, device=dev, dtype=dtype), |
| 53 | A_log=torch.randn(NUM_HEADS, device=dev, dtype=torch.float32) * 0.5, |
| 54 | dt_bias=torch.randn(NUM_HEADS * HEAD_DIM, device=dev, dtype=torch.float32) |
| 55 | * 0.5, |
| 56 | h0=torch.randn( |
| 57 | len(seqlens), NUM_HEADS, HEAD_DIM, HEAD_DIM, device=dev, dtype=torch.float32 |
| 58 | ) |
| 59 | * 0.1, |
| 60 | cu=cu, |
| 61 | ) |
| 62 | |
| 63 | |
| 64 | def _run(inp: dict, use_fused: bool, out: torch.Tensor | None = None, **kw): |
| 65 | from vllm.models.kimi_k3.amd.ops.kda_prefill import chunk_kda_prefill |
| 66 | |
| 67 | return chunk_kda_prefill( |
| 68 | q=inp["q"].clone(), |
| 69 | k=inp["k"].clone(), |
| 70 | v=inp["v"].clone(), |
| 71 | raw_g=inp["raw_g"].clone(), |
| 72 | raw_beta=inp["raw_beta"], |
| 73 | A_log=inp["A_log"], |
| 74 | g_bias=inp["dt_bias"], |
| 75 | lower_bound=LOWER_BOUND, |
| 76 | initial_state=inp["h0"].clone(), |
| 77 | output_final_state=True, |
| 78 | use_qk_l2norm_in_kernel=True, |
| 79 | cu_seqlens=inp["cu"], |
| 80 | use_fused_chunk=use_fused, |
| 81 | out=out, |
| 82 | **kw, |
| 83 | ) |
| 84 | |
| 85 | |
| 86 | def _force_groups(monkeypatch, groups: int) -> None: |
| 87 | """Pin the group count, which is otherwise a function of the device.""" |
| 88 | import vllm.models.kimi_k3.amd.ops.kda_chunk as kda_chunk |
| 89 | |
| 90 | monkeypatch.setattr(kda_chunk, "_chunk_groups", lambda *_: groups) |
| 91 | |
| 92 | |
| 93 | @pytest.mark.parametrize( |
| 94 | "seqlens", |
| 95 | [ |
| 96 | [1], # shorter than a chunk |
| 97 | [64], # exactly one chunk |
| 98 | [1024], |
| 99 | [4096], # many chunks, one sequence |
| 100 | [513, 64, 1, 1200], # ragged, including a partial tail chunk |
| 101 | [512] * 6, |
| 102 | ], |
| 103 | ) |
| 104 | def test_fused_chunk_matches_triton(seqlens: list[int]) -> None: |
| 105 | _requires_kernel() |
| 106 | |
| 107 | inp = _inputs(seqlens, seed=len(seqlens)) |
| 108 | o_ref, state_ref = _run(inp, use_fused=False) |
| 109 | o_fused, state_fused = _run(inp, use_fused=True) |
| 110 | |
| 111 | torch.testing.assert_close(o_fused.float(), o_ref.float(), rtol=2e-2, atol=2e-2) |
| 112 | torch.testing.assert_close(state_fused, state_ref, rtol=1e-3, atol=1e-3) |
| 113 | |
| 114 | |
| 115 | @pytest.mark.parametrize("use_fused", [False, True]) |
| 116 | def test_out_argument_receives_the_result(use_fused: bool) -> None: |
| 117 | """``out`` is a promise about where the result lands. |
| 118 | |
| 119 | The layer passes a slice of its own output buffer so neither path has to |
| 120 | copy afterwards, and it then skips the copy based on the returned tensor |
| 121 | aliasing that buffer. Both must hold for the fused path and for the Triton |
| 122 | fallback, which allocates its own buffer internally. |
| 123 | """ |
| 124 | if use_fused: |
| 125 | _requires_kernel() |
| 126 | |
| 127 | seqlens = [321, 64, 700] |
| 128 | inp = _inputs(seqlens, seed=13) |
| 129 | o_ref, _ = _run(inp, use_fused=use_fused) |
| 130 | |
| 131 | # A non-zero offset into a larger buffer, as the layer produces when a |
| 132 | # mixed batch puts decode tokens first. |
| 133 | total = sum(seqlens) |
| 134 | buf = torch.empty( |
| 135 | 1, total + 8, NUM_HEADS, HEAD_DIM, device="cuda", dtype=o_ref.dtype |
| 136 | ) |
| 137 | view = buf[:, 8:] |
| 138 | o_got, _ = _run(inp, use_fused=use_fused, out=view) |
| 139 | |
| 140 | assert o_got.data_ptr() == view.data_ptr() |
| 141 | torch.testing.assert_close(view.float(), o_ref.float(), rtol=2e-2, atol=2e-2) |
| 142 | |
| 143 | |
| 144 | def test_fused_chunk_accepts_projection_views() -> None: |
| 145 | """`beta` and `g` arrive as last-dim slices of the fused QKVGFAB projection. |
| 146 | |
| 147 | Those views carry the projection's row stride, so a kernel that indexes them |
| 148 | densely reads the wrong elements or rejects the tensor outright. Build the |
| 149 | inputs the way the layer does rather than as standalone tensors. |
| 150 | """ |
| 151 | _requires_kernel() |
| 152 | |
| 153 | seqlens = [777, 64, 391] |
| 154 | inp = _inputs(seqlens, seed=7) |
| 155 | total = sum(seqlens) |
| 156 | |
| 157 | # One wide projection, split on the last dim exactly like KimiK3DeltaAttention. |
| 158 | qkv_width = 3 * NUM_HEADS * HEAD_DIM |
| 159 | projection = torch.randn( |
| 160 | total, qkv_width + NUM_HEADS, device="cuda", dtype=torch.bfloat16 |
| 161 | ) |
| 162 | beta_view = projection.split([qkv_width, NUM_HEADS], dim=-1)[1].unsqueeze(0) |
| 163 | assert not beta_view.is_contiguous() |
| 164 | |
| 165 | inp["raw_beta"] = beta_view.contiguous() |
| 166 | o_ref, state_ref = _run(inp, use_fused=False) |
| 167 | |
| 168 | inp["raw_beta"] = beta_view |
| 169 | o_fused, state_fused = _run(inp, use_fused=True) |
| 170 | |
| 171 | torch.testing.assert_close(o_fused.float(), o_ref.float(), rtol=2e-2, atol=2e-2) |
| 172 | torch.testing.assert_close(state_fused, state_ref, rtol=1e-3, atol=1e-3) |
| 173 | |
| 174 | |
| 175 | def test_fused_chunk_survives_correlated_keys_and_real_gate_bias() -> None: |
| 176 | """The input regime the served model actually produces. |
| 177 | |
| 178 | This is what forces the intra-chunk inverse to keep fp32 MFMA operands: at |
| 179 | bf16 the Neumann doubling's intermediate powers cancel catastrophically |
| 180 | here, diverging by orders of magnitude and driving the final state to inf. |
| 181 | |
| 182 | Three properties of real traffic combine here and in none of the tests |
| 183 | above, all measured from tensors captured off a live TP8 server: |
| 184 | |
| 185 | * repetitive text gives near-duplicate L2-normalized keys, so ``k_i . k_j`` |
| 186 | approaches 1 and the Neumann terms stop shrinking; |
| 187 | * the checkpoint's ``dt_bias`` lies in ``[-7.8, -1.4]``, which makes the |
| 188 | within-chunk gate cumsum span ~300 powers of two; |
| 189 | * ``beta`` saturates near 0.95. |
| 190 | |
| 191 | ``_inputs`` draws all three from ``randn``, landing in the benign regime. |
| 192 | """ |
| 193 | _requires_kernel() |
| 194 | |
| 195 | torch.manual_seed(11) |
| 196 | dev, dtype = "cuda", torch.bfloat16 |
| 197 | total = 1088 |
| 198 | shape = (1, total, NUM_HEADS, HEAD_DIM) |
| 199 | base = torch.randn(1, 1, NUM_HEADS, HEAD_DIM, device=dev, dtype=dtype) |
| 200 | |
| 201 | def correlated(): |
| 202 | return (base + torch.randn(shape, device=dev, dtype=dtype) * 0.02).contiguous() |
| 203 | |
| 204 | inp = dict( |
| 205 | q=correlated(), |
| 206 | k=correlated(), |
| 207 | v=torch.randn(shape, device=dev, dtype=dtype) * 0.5, |
| 208 | raw_g=torch.randn(shape, device=dev, dtype=dtype), |
| 209 | raw_beta=torch.full((1, total, NUM_HEADS), 3.0, device=dev, dtype=dtype), |
| 210 | A_log=torch.rand(NUM_HEADS, device=dev, dtype=torch.float32) * 0.6 - 0.42, |
| 211 | dt_bias=torch.rand(NUM_HEADS * HEAD_DIM, device=dev, dtype=torch.float32) * 6.4 |
| 212 | - 7.8, |
| 213 | h0=torch.zeros( |
| 214 | 1, NUM_HEADS, HEAD_DIM, HEAD_DIM, device=dev, dtype=torch.float32 |
| 215 | ), |
| 216 | cu=torch.tensor([0, total], device=dev, dtype=torch.int32), |
| 217 | ) |
| 218 | |
| 219 | o_ref, state_ref = _run(inp, use_fused=False) |
| 220 | o_fused, state_fused = _run(inp, use_fused=True) |
| 221 | |
| 222 | # Looser than the tests above: near-identical keys plus a saturated beta |
| 223 | # are past anything real text produces, and the inverse is rounded to bf16 |
| 224 | # before the w/u GEMMs. The tolerance still catches divergence, which in |
| 225 | # this regime is orders of magnitude rather than a few times the bound. |
| 226 | torch.testing.assert_close(o_fused.float(), o_ref.float(), rtol=5e-2, atol=5e-2) |
| 227 | torch.testing.assert_close(state_fused, state_ref, rtol=5e-2, atol=5e-2) |
| 228 | |
| 229 | |
| 230 | @pytest.mark.parametrize( |
| 231 | "seqlens", |
| 232 | [ |
| 233 | [1], # every conv tap comes from the cache |
| 234 | [2], # cache and batch both contribute |
| 235 | [64], |
| 236 | [1024], |
| 237 | [513, 64, 1, 1200], # ragged, tail chunk shorter than the conv width |
| 238 | ], |
| 239 | ) |
| 240 | @pytest.mark.parametrize("has_state", [False, True]) |
| 241 | def test_fused_conv_matches_reference_conv(seqlens: list[int], has_state: bool) -> None: |
| 242 | """The prologue's optional fused convolution must match a separate conv. |
| 243 | |
| 244 | The layer does not use this path -- recomputing the conv once per load |
| 245 | phase costs more than the ``causal_conv1d_fn`` launches it saves -- so it |
| 246 | is driven through the ops directly rather than through |
| 247 | ``chunk_kda_prefill``, which does not expose it. |
| 248 | |
| 249 | The reference is computed in torch rather than with ``causal_conv1d_fn``, |
| 250 | which is not self-consistent for these shapes on ROCm (it leaves part of |
| 251 | its output unwritten, so repeated identical calls disagree). |
| 252 | |
| 253 | The cases that matter are the ones touching the sequence boundary, where |
| 254 | the fused path takes taps from the conv cache instead of the batch: |
| 255 | sequences shorter than the kernel width, and a cache that is or is not |
| 256 | populated. The rolled-forward cache is checked too, since a wrong shift |
| 257 | only shows up on the next step. |
| 258 | """ |
| 259 | import torch.nn.functional as F |
| 260 | |
| 261 | from vllm.models.kimi_k3.amd.ops.kda_chunk import ( |
| 262 | fused_kda_chunk, |
| 263 | fused_kda_prologue, |
| 264 | ) |
| 265 | |
| 266 | _requires_kernel() |
| 267 | |
| 268 | dev, dtype = "cuda", torch.bfloat16 |
| 269 | lp = NUM_HEADS * HEAD_DIM |
| 270 | width, n = 4, len(seqlens) |
| 271 | inp = _inputs(seqlens, seed=7 + len(seqlens) + int(has_state)) |
| 272 | total = int(inp["cu"][-1]) |
| 273 | |
| 274 | torch.manual_seed(11 + int(has_state)) |
| 275 | # The layer hands the prologue a strided view of the QKVGFAB projection. |
| 276 | projection = torch.randn(total, 3 * lp + 64, device=dev, dtype=dtype) * 0.5 |
| 277 | mixed = projection[:, : 3 * lp] |
| 278 | # conv1d is built with bias=False and fp32 weights; the fused kernel takes |
| 279 | # the width-major [qkv, width, channel] mirror the decode path already |
| 280 | # keeps. |
| 281 | weight = torch.randn(3, width, lp, device=dev, dtype=torch.float32) * 0.3 |
| 282 | store = torch.zeros(n, width - 1, 3 * lp, device=dev, dtype=dtype) |
| 283 | if has_state: |
| 284 | store.normal_(0.0, 0.5) |
| 285 | state = store.transpose(-1, -2) |
| 286 | idx = torch.arange(n, device=dev, dtype=torch.int32) |
| 287 | has_init = torch.full((n,), has_state, device=dev, dtype=torch.bool) |
| 288 | |
| 289 | # Reference conv, per sequence, in fp32 then rounded like a bf16 output. |
| 290 | conv_ref = torch.empty(total, 3 * lp, device=dev, dtype=dtype) |
| 291 | state_ref = torch.empty_like(store) |
| 292 | flat_w = weight.permute(0, 2, 1).reshape(3 * lp, width) |
| 293 | for i, (b, e) in enumerate(zip(inp["cu"][:-1].tolist(), inp["cu"][1:].tolist())): |
| 294 | pre = ( |
| 295 | store[i].float() |
| 296 | if has_state |
| 297 | else torch.zeros(width - 1, 3 * lp, device=dev, dtype=torch.float32) |
| 298 | ) |
| 299 | xs = torch.cat([pre, mixed[b:e].float()], 0) |
| 300 | acc = torch.zeros(e - b, 3 * lp, device=dev, dtype=torch.float32) |
| 301 | for w in range(width): |
| 302 | acc += xs[w : w + (e - b)] * flat_w[:, w].unsqueeze(0) |
| 303 | conv_ref[b:e] = F.silu(acc).to(dtype) |
| 304 | state_ref[i] = xs[-(width - 1) :].to(dtype) |
| 305 | |
| 306 | def _bands(x, contiguous): |
| 307 | out = [ |
| 308 | b.unsqueeze(0).unflatten(-1, (NUM_HEADS, HEAD_DIM)) |
| 309 | for b in x.split(lp, dim=-1) |
| 310 | ] |
| 311 | return [b.contiguous() for b in out] if contiguous else out |
| 312 | |
| 313 | # Compare through both kernels rather than the prologue workspaces: `aqk` |
| 314 | # is deliberately left unwritten past each chunk's valid rows, so those |
| 315 | # elements are whatever the allocator last held. |
| 316 | scale = HEAD_DIM**-0.5 |
| 317 | |
| 318 | def _through_kernels(q, k, v, conv: dict | None): |
| 319 | ws = fused_kda_prologue( |
| 320 | q=q, |
| 321 | k=k, |
| 322 | v=v, |
| 323 | raw_g=inp["raw_g"], |
| 324 | raw_beta=inp["raw_beta"], |
| 325 | A_log=inp["A_log"], |
| 326 | dt_bias=inp["dt_bias"], |
| 327 | scale=scale, |
| 328 | lower_bound=LOWER_BOUND, |
| 329 | cu_seqlens=inp["cu"], |
| 330 | **(conv or {}), |
| 331 | ) |
| 332 | return fused_kda_chunk( |
| 333 | qg=ws["qg"], |
| 334 | w=ws["w"], |
| 335 | u=ws["u"], |
| 336 | kg_t=ws["kg_t"], |
| 337 | aqk=ws["aqk"], |
| 338 | decay=ws["decay"], |
| 339 | out=torch.empty_like(ws["u"]), |
| 340 | scale=scale, |
| 341 | cu_seqlens=inp["cu"], |
| 342 | initial_state=inp["h0"].clone(), |
| 343 | output_final_state=True, |
| 344 | ) |
| 345 | |
| 346 | qr, kr, vr = _bands(conv_ref, True) |
| 347 | o_ref, s_ref = _through_kernels(qr, kr, vr, None) |
| 348 | |
| 349 | got_state = state.clone() |
| 350 | qf, kf, vf = _bands(mixed, False) |
| 351 | o_got, s_got = _through_kernels( |
| 352 | qf, |
| 353 | kf, |
| 354 | vf, |
| 355 | dict( |
| 356 | conv_weight=weight, |
| 357 | conv_state=got_state, |
| 358 | conv_state_indices=idx, |
| 359 | conv_has_initial_state=has_init, |
| 360 | ), |
| 361 | ) |
| 362 | |
| 363 | torch.testing.assert_close(o_got.float(), o_ref.float(), rtol=2e-2, atol=2e-2) |
| 364 | torch.testing.assert_close(s_got, s_ref, rtol=1e-3, atol=1e-3) |
| 365 | torch.testing.assert_close( |
| 366 | got_state.transpose(-1, -2).float(), state_ref.float(), rtol=0, atol=0 |
| 367 | ) |
| 368 | |
| 369 | |
| 370 | @pytest.mark.parametrize("groups", [2, 4, 8]) |
| 371 | @pytest.mark.parametrize("seqlens", [[64], [64, 64], [100, 64], [513, 64, 1]]) |
| 372 | def test_chunk_group_scan_matches_serial_walk( |
| 373 | groups: int, seqlens: list[int], monkeypatch |
| 374 | ) -> None: |
| 375 | """Splitting the chunk walk into parallel groups must not change the result. |
| 376 | |
| 377 | Two bugs lived here, both invisible while the path was gated to >= 128 |
| 378 | chunks per sequence: |
| 379 | |
| 380 | * a group with no chunks returned before writing its ``B_g``, so the scan |
| 381 | folded whatever the workspace last held into the final state — these |
| 382 | shapes have fewer chunks than groups, so most groups are empty; |
| 383 | * the second pass started every group after the first from a zero state |
| 384 | instead of the one the scan composed, which only looked right because a |
| 385 | strongly decaying gate makes the dropped term small. |
| 386 | """ |
| 387 | _requires_kernel() |
| 388 | |
| 389 | inp = _inputs(seqlens, seed=5 + groups) |
| 390 | _force_groups(monkeypatch, 1) |
| 391 | o_ref, state_ref = _run(inp, use_fused=True) |
| 392 | _force_groups(monkeypatch, groups) |
| 393 | o_got, state_got = _run(inp, use_fused=True) |
| 394 | |
| 395 | torch.testing.assert_close(o_got.float(), o_ref.float(), rtol=2e-2, atol=2e-2) |
| 396 | torch.testing.assert_close(state_got, state_ref, rtol=1e-3, atol=1e-3) |
| 397 | |
| 398 | |
| 399 | @pytest.mark.parametrize("seqlens", [[320], [513, 64, 1, 1200], [64, 64, 64]]) |
| 400 | def test_accepts_chunk_metadata_from_the_attention_builder( |
| 401 | seqlens: list[int], |
| 402 | ) -> None: |
| 403 | """The layer passes precomputed chunk metadata, not None. |
| 404 | |
| 405 | ``KimiK3ROCmKDAMetadataBuilder`` builds these on device to avoid a blocking |
| 406 | device->host copy. Passing ``None`` instead takes a different path that |
| 407 | builds them locally, so drive the op with the builder's own tensors. |
| 408 | """ |
| 409 | _requires_kernel() |
| 410 | |
| 411 | from vllm.models.kimi_k3.amd.kda_metadata import prepare_chunk_metadata_device |
| 412 | from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE |
| 413 | |
| 414 | inp = _inputs(seqlens, seed=17) |
| 415 | cu = inp["cu"] |
| 416 | chunk_indices, chunk_offsets = prepare_chunk_metadata_device( |
| 417 | cu, cu.cpu(), FLA_CHUNK_SIZE |
| 418 | ) |
| 419 | |
| 420 | o_ref, state_ref = _run(inp, use_fused=True) |
| 421 | o_got, state_got = _run( |
| 422 | inp, use_fused=True, chunk_indices=chunk_indices, chunk_offsets=chunk_offsets |
| 423 | ) |
| 424 | |
| 425 | torch.testing.assert_close(o_got.float(), o_ref.float(), rtol=2e-2, atol=2e-2) |
| 426 | torch.testing.assert_close(state_got, state_ref, rtol=1e-3, atol=1e-3) |
| 427 | |
| 428 | |
| 429 | def _checkpoint_buffer(rows: int) -> torch.Tensor: |
| 430 | """A destination filled with NaN, so an unwritten row is unmistakable.""" |
| 431 | return torch.full( |
| 432 | (rows, NUM_HEADS, HEAD_DIM, HEAD_DIM), |
| 433 | float("nan"), |
| 434 | device="cuda", |
| 435 | dtype=torch.float32, |
| 436 | ) |
| 437 | |
| 438 | |
| 439 | def _prefix_inputs(inp: dict, seq: int, num_tokens: int) -> dict: |
| 440 | """Sequence ``seq`` truncated to its first ``num_tokens`` tokens.""" |
| 441 | bos = int(inp["cu"][seq]) |
| 442 | tokens = slice(bos, bos + num_tokens) |
| 443 | prefix = dict(inp) |
| 444 | for key in ("q", "k", "v", "raw_g", "raw_beta"): |
| 445 | prefix[key] = inp[key][:, tokens].contiguous() |
| 446 | prefix["h0"] = inp["h0"][seq : seq + 1].contiguous() |
| 447 | prefix["cu"] = torch.tensor([0, num_tokens], device="cuda", dtype=torch.int32) |
| 448 | return prefix |
| 449 | |
| 450 | |
| 451 | def _run_with_checkpoints( |
| 452 | inp: dict, offsets: list[int], rows: int | None = None, **kw |
| 453 | ) -> tuple[torch.Tensor, tuple]: |
| 454 | ckpt = _checkpoint_buffer(len(offsets) if rows is None else rows) |
| 455 | result = _run( |
| 456 | inp, |
| 457 | use_fused=True, |
| 458 | checkpoint_state=ckpt, |
| 459 | checkpoint_offsets=torch.tensor(offsets, device="cuda", dtype=torch.int32), |
| 460 | **kw, |
| 461 | ) |
| 462 | return ckpt, result |
| 463 | |
| 464 | |
| 465 | @pytest.mark.parametrize( |
| 466 | "seqlens,offsets", |
| 467 | [ |
| 468 | ([256], [64]), # the first interior boundary |
| 469 | ([256], [192]), # the last one |
| 470 | ([1024], [512]), |
| 471 | ([320, 320], [64, 256]), # every sequence checkpoints |
| 472 | ([513, 64, 1, 1200], [448, 0, 0, 640]), # ragged, most opted out |
| 473 | ], |
| 474 | ) |
| 475 | def test_checkpoint_matches_a_prefill_that_stops_at_the_offset( |
| 476 | seqlens: list[int], offsets: list[int], monkeypatch |
| 477 | ) -> None: |
| 478 | """A checkpoint is only useful if resuming from it is indistinguishable. |
| 479 | |
| 480 | The scheduler serves the exported state on a later prefix-cache hit as if |
| 481 | the prompt had been prefilled up to that block and no further, so the |
| 482 | kernel has to hand back exactly what a prefill truncated there produces -- |
| 483 | not merely something close. The walk passes through that state anyway, so |
| 484 | equality is bit-exact and the test asserts it as such. |
| 485 | |
| 486 | Sequences with offset 0 opt out, and their rows must stay untouched. |
| 487 | """ |
| 488 | _requires_kernel() |
| 489 | _force_groups(monkeypatch, 1) |
| 490 | |
| 491 | inp = _inputs(seqlens, seed=23 + len(seqlens)) |
| 492 | ckpt, _ = _run_with_checkpoints(inp, offsets) |
| 493 | |
| 494 | for i, offset in enumerate(offsets): |
| 495 | if offset == 0: |
| 496 | assert ckpt[i].isnan().all() |
| 497 | continue |
| 498 | _, truncated = _run(_prefix_inputs(inp, i, offset), use_fused=True) |
| 499 | torch.testing.assert_close(ckpt[i], truncated[0], rtol=0, atol=0) |
| 500 | |
| 501 | |
| 502 | def test_checkpoint_leaves_the_output_and_final_state_alone() -> None: |
| 503 | """Exporting a checkpoint must not perturb the walk it is taken from.""" |
| 504 | _requires_kernel() |
| 505 | |
| 506 | seqlens = [768, 320] |
| 507 | inp = _inputs(seqlens, seed=29) |
| 508 | o_ref, state_ref = _run(inp, use_fused=True) |
| 509 | _, (o_got, state_got) = _run_with_checkpoints(inp, [512, 128]) |
| 510 | |
| 511 | torch.testing.assert_close(o_got.float(), o_ref.float(), rtol=0, atol=0) |
| 512 | torch.testing.assert_close(state_got, state_ref, rtol=0, atol=0) |
| 513 | |
| 514 | |
| 515 | @pytest.mark.parametrize("offset", [100, 63, 65, 1024]) |
| 516 | def test_unusable_checkpoint_offsets_write_nothing(offset: int) -> None: |
| 517 | """Offsets the walk never reaches must leave the destination untouched. |
| 518 | |
| 519 | The recurrent state only exists at chunk boundaries, so an offset that is |
| 520 | not a multiple of ``KDA_CHECKPOINT_ALIGNMENT`` -- or that lies past the end |
| 521 | of the sequence -- cannot be honoured. Rounding it to a nearby boundary |
| 522 | would hand the scheduler a state for the wrong token count, so the kernel |
| 523 | stores nothing and leaves the caller to filter these out before the block |
| 524 | is hashed. |
| 525 | """ |
| 526 | _requires_kernel() |
| 527 | |
| 528 | from vllm.models.kimi_k3.amd.ops.kda_chunk import KDA_CHECKPOINT_ALIGNMENT |
| 529 | |
| 530 | assert offset % KDA_CHECKPOINT_ALIGNMENT != 0 or offset > 512 |
| 531 | inp = _inputs([512], seed=31) |
| 532 | ckpt, _ = _run_with_checkpoints(inp, [offset]) |
| 533 | |
| 534 | assert ckpt.isnan().all() |
| 535 | |
| 536 | |
| 537 | def test_checkpoint_state_indices_pick_the_destination_row() -> None: |
| 538 | """With indices the export lands straight in the paged state cache. |
| 539 | |
| 540 | FlashKDA writes one row per sequence into a staging buffer that the caller |
| 541 | then scatters. The ROCm walk holds the state in registers already, so it |
| 542 | can be told which cache block each sequence owns and skip the staging |
| 543 | buffer and its scatter launch entirely. A negative row opts a sequence out, |
| 544 | the same way a zero offset does. |
| 545 | """ |
| 546 | _requires_kernel() |
| 547 | |
| 548 | seqlens = [256, 256, 256] |
| 549 | rows = 8 |
| 550 | inp = _inputs(seqlens, seed=37) |
| 551 | indices = [5, -1, 2] |
| 552 | ckpt, _ = _run_with_checkpoints( |
| 553 | inp, |
| 554 | [64, 64, 128], |
| 555 | rows=rows, |
| 556 | checkpoint_state_indices=torch.tensor( |
| 557 | indices, device="cuda", dtype=torch.int32 |
| 558 | ), |
| 559 | ) |
| 560 | |
| 561 | written = {5: (0, 64), 2: (2, 128)} |
| 562 | for row in range(rows): |
| 563 | if row not in written: |
| 564 | assert ckpt[row].isnan().all(), f"row {row} was written" |
| 565 | continue |
| 566 | seq, offset = written[row] |
| 567 | _, truncated = _run(_prefix_inputs(inp, seq, offset), use_fused=True) |
| 568 | torch.testing.assert_close(ckpt[row], truncated[0], rtol=0, atol=0) |
| 569 | |
| 570 | |
| 571 | @pytest.mark.parametrize("groups", [2, 4]) |
| 572 | def test_checkpoint_survives_the_group_split(groups: int, monkeypatch) -> None: |
| 573 | """The split walks each group from a composed state, not from zero. |
| 574 | |
| 575 | Only the second pass holds the true state at an interior boundary, so a |
| 576 | checkpoint taken from the first pass would be the group-local state and |
| 577 | silently wrong. |
| 578 | """ |
| 579 | _requires_kernel() |
| 580 | |
| 581 | inp = _inputs([1024], seed=41) |
| 582 | _force_groups(monkeypatch, 1) |
| 583 | _, serial = _run(_prefix_inputs(inp, 0, 512), use_fused=True) |
| 584 | _force_groups(monkeypatch, groups) |
| 585 | ckpt, _ = _run_with_checkpoints(inp, [512]) |
| 586 | |
| 587 | torch.testing.assert_close(ckpt[0], serial[0], rtol=1e-3, atol=1e-3) |
| 588 | |
| 589 | |
| 590 | def test_triton_fallback_refuses_a_checkpoint_export() -> None: |
| 591 | """Dropping the export silently would serve an unwritten state. |
| 592 | |
| 593 | The scheduler allocates and hashes the checkpoint block on the strength of |
| 594 | the layer opting in, so a backend that cannot write it has to say so rather |
| 595 | than return a plausible output. |
| 596 | """ |
| 597 | inp = _inputs([256], seed=43) |
| 598 | with pytest.raises(NotImplementedError): |
| 599 | _run( |
| 600 | inp, |
| 601 | use_fused=False, |
| 602 | checkpoint_state=_checkpoint_buffer(1), |
| 603 | checkpoint_offsets=torch.tensor([64], device="cuda", dtype=torch.int32), |
| 604 | ) |
| 605 | |
| 606 | |
| 607 | # --------------------------------------------------------------------------- |
| 608 | # Paged recurrent state |
| 609 | # |
| 610 | # Each sequence reads its initial state from a row of the paged cache and |
| 611 | # writes its final state back to the same row. These check that the result |
| 612 | # matches the dense initial_state/final_state pair, that a sequence marked |
| 613 | # cold starts from zeros rather than from whatever its row held, and that rows |
| 614 | # no sequence owns are left untouched. |
| 615 | # --------------------------------------------------------------------------- |
| 616 | |
| 617 | |
| 618 | def _paged_cache(inp: dict, warm: list[bool], rows: list[int], slots: int): |
| 619 | """A cache holding each sequence's h0 at its row, a fill pattern elsewhere.""" |
| 620 | cache = torch.full( |
| 621 | (slots, NUM_HEADS, HEAD_DIM, HEAD_DIM), |
| 622 | -7.0, |
| 623 | device="cuda", |
| 624 | dtype=torch.float32, |
| 625 | ) |
| 626 | idx = torch.tensor(rows, device="cuda", dtype=torch.int32) |
| 627 | for n, row in enumerate(rows): |
| 628 | # A cold row keeps its fill pattern, which must not reach the walk. |
| 629 | if warm[n]: |
| 630 | cache[row] = inp["h0"][n] |
| 631 | return cache, idx, torch.tensor(warm, device="cuda", dtype=torch.bool) |
| 632 | |
| 633 | |
| 634 | def _run_paged(inp: dict, cache, idx, warm, use_fused: bool, **kw): |
| 635 | from vllm.models.kimi_k3.amd.ops.kda_prefill import chunk_kda_prefill |
| 636 | |
| 637 | return chunk_kda_prefill( |
| 638 | q=inp["q"].clone(), |
| 639 | k=inp["k"].clone(), |
| 640 | v=inp["v"].clone(), |
| 641 | raw_g=inp["raw_g"].clone(), |
| 642 | raw_beta=inp["raw_beta"], |
| 643 | A_log=inp["A_log"], |
| 644 | g_bias=inp["dt_bias"], |
| 645 | lower_bound=LOWER_BOUND, |
| 646 | state_cache=cache, |
| 647 | state_indices=idx, |
| 648 | has_initial_state=warm, |
| 649 | use_qk_l2norm_in_kernel=True, |
| 650 | cu_seqlens=inp["cu"], |
| 651 | use_fused_chunk=use_fused, |
| 652 | **kw, |
| 653 | ) |
| 654 | |
| 655 | |
| 656 | @pytest.mark.parametrize("use_fused", [False, True]) |
| 657 | @pytest.mark.parametrize("seqlens", [[512], [2048], [320, 64], [513, 64, 1]]) |
| 658 | def test_paged_state_matches_the_dense_pair( |
| 659 | use_fused: bool, seqlens: list[int] |
| 660 | ) -> None: |
| 661 | """A paged cache must give the same result as the dense pair.""" |
| 662 | _requires_kernel() |
| 663 | inp = _inputs(seqlens, seed=11) |
| 664 | n = len(seqlens) |
| 665 | rows = [3 * i + 5 for i in range(n)] |
| 666 | cache, idx, warm = _paged_cache(inp, [True] * n, rows, slots=3 * n + 12) |
| 667 | before = cache.clone() |
| 668 | |
| 669 | o_ref, ht_ref = _run(inp, use_fused=use_fused) |
| 670 | o_got, ht_got = _run_paged(inp, cache, idx, warm, use_fused=use_fused) |
| 671 | |
| 672 | assert ht_got is None, "the final state is in the cache, not returned" |
| 673 | torch.testing.assert_close(o_got, o_ref, rtol=5e-3, atol=5e-3) |
| 674 | torch.testing.assert_close(cache[idx.long()], ht_ref, rtol=5e-3, atol=5e-3) |
| 675 | |
| 676 | untouched = torch.ones(cache.shape[0], dtype=torch.bool, device="cuda") |
| 677 | untouched[idx.long()] = False |
| 678 | assert torch.equal(cache[untouched], before[untouched]), ( |
| 679 | "a row no sequence owns was written" |
| 680 | ) |
| 681 | |
| 682 | |
| 683 | @pytest.mark.parametrize("use_fused", [False, True]) |
| 684 | def test_cold_rows_start_from_zero_not_from_cache_junk(use_fused: bool) -> None: |
| 685 | """`has_initial_state=False` must ignore whatever the row already holds.""" |
| 686 | _requires_kernel() |
| 687 | seqlens = [256, 256, 256] |
| 688 | inp = _inputs(seqlens, seed=12) |
| 689 | warm = [True, False, True] |
| 690 | rows = [4, 9, 1] |
| 691 | cache, idx, warm_t = _paged_cache(inp, warm, rows, slots=16) |
| 692 | |
| 693 | # The dense reference zeroes exactly the cold sequences' initial states. |
| 694 | ref = dict(inp) |
| 695 | ref["h0"] = inp["h0"].clone() |
| 696 | ref["h0"][1] = 0.0 |
| 697 | o_ref, ht_ref = _run(ref, use_fused=use_fused) |
| 698 | o_got, _ = _run_paged(inp, cache, idx, warm_t, use_fused=use_fused) |
| 699 | |
| 700 | torch.testing.assert_close(o_got, o_ref, rtol=5e-3, atol=5e-3) |
| 701 | torch.testing.assert_close(cache[idx.long()], ht_ref, rtol=5e-3, atol=5e-3) |
| 702 | |
| 703 | |
| 704 | @pytest.mark.parametrize("use_fused", [False, True]) |
| 705 | def test_state_indices_accept_a_block_table_column(use_fused: bool) -> None: |
| 706 | # The layer passes block_table[:, 0], which is contiguous only while the |
| 707 | # table has one column. Speculative decoding reserves a state page per |
| 708 | # draft token, so the table widens and that column becomes strided. |
| 709 | _requires_kernel() |
| 710 | seqlens = [512, 320, 64] |
| 711 | inp = _inputs(seqlens, seed=43) |
| 712 | n = len(seqlens) |
| 713 | rows = [3 * i + 5 for i in range(n)] |
| 714 | cache, idx, warm = _paged_cache(inp, [True] * n, rows, slots=3 * n + 12) |
| 715 | ref_cache = cache.clone() |
| 716 | |
| 717 | table = torch.zeros(n, 9, dtype=torch.int32, device="cuda") |
| 718 | table[:, 0] = idx |
| 719 | column = table[:, 0] |
| 720 | assert not column.is_contiguous() |
| 721 | |
| 722 | o_ref, _ = _run_paged(inp, ref_cache, idx, warm, use_fused=use_fused) |
| 723 | o_got, _ = _run_paged(inp, cache, column, warm, use_fused=use_fused) |
| 724 | |
| 725 | torch.testing.assert_close(o_got, o_ref) |
| 726 | torch.testing.assert_close(cache, ref_cache) |
| 727 | |
| 728 | |
| 729 | @pytest.mark.parametrize("use_fused", [False, True]) |
| 730 | def test_padded_zero_length_sequence_is_inert(use_fused: bool) -> None: |
| 731 | # A batch carrying speculative decodes pads its non-spec group with |
| 732 | # zero-length sequences whose cache row still belongs to a request, so |
| 733 | # consuming no tokens has to mean writing nothing and perturbing nobody. |
| 734 | _requires_kernel() |
| 735 | live = [200, 1, 128] |
| 736 | # _inputs draws h0 last, so the padded batch's live rows are the same draw |
| 737 | # as the unpadded one's and the two runs compare elementwise. |
| 738 | padded, unpadded = _inputs([*live, 0], seed=41), _inputs(live, seed=41) |
| 739 | rows = [3 * i + 5 for i in range(4)] |
| 740 | |
| 741 | cache, idx, warm = _paged_cache(padded, [True] * 4, rows, slots=24) |
| 742 | before = cache.clone() |
| 743 | _run_paged(padded, cache, idx, warm, use_fused=use_fused) |
| 744 | |
| 745 | ref_cache, ref_idx, ref_warm = _paged_cache(unpadded, [True] * 3, rows[:3], 24) |
| 746 | _run_paged(unpadded, ref_cache, ref_idx, ref_warm, use_fused=use_fused) |
| 747 | |
| 748 | assert torch.equal(cache[rows[3]], before[rows[3]]), ( |
| 749 | "the padded sequence's cache row was written" |
| 750 | ) |
| 751 | torch.testing.assert_close(cache[ref_idx.long()], ref_cache[ref_idx.long()]) |
| 752 | |
| 753 | |
| 754 | @pytest.mark.parametrize("groups", [2, 4]) |
| 755 | def test_paged_state_survives_the_group_split(groups: int, monkeypatch) -> None: |
| 756 | """At G>1 the read and the write are different blocks of one launch.""" |
| 757 | _requires_kernel() |
| 758 | _force_groups(monkeypatch, groups) |
| 759 | inp = _inputs([1024, 512], seed=13) |
| 760 | rows = [6, 2] |
| 761 | cache, idx, warm = _paged_cache(inp, [True, True], rows, slots=12) |
| 762 | |
| 763 | o_ref, ht_ref = _run(inp, use_fused=True) |
| 764 | o_got, _ = _run_paged(inp, cache, idx, warm, use_fused=True) |
| 765 | |
| 766 | torch.testing.assert_close(o_got, o_ref, rtol=5e-3, atol=5e-3) |
| 767 | torch.testing.assert_close(cache[idx.long()], ht_ref, rtol=5e-3, atol=5e-3) |
| 768 | |
| 769 | |
| 770 | def test_paged_state_and_checkpoint_export_together() -> None: |
| 771 | """The checkpoint lands in its own row while the walk owns the state row.""" |
| 772 | _requires_kernel() |
| 773 | inp = _inputs([1024], seed=14) |
| 774 | cache, idx, warm = _paged_cache(inp, [True], [3], slots=12) |
| 775 | ckidx = torch.tensor([8], device="cuda", dtype=torch.int32) |
| 776 | |
| 777 | o_ref, ht_ref = _run(inp, use_fused=True) |
| 778 | prefix = _prefix_inputs(inp, 0, 512) |
| 779 | _, ck_ref = _run(prefix, use_fused=True) |
| 780 | |
| 781 | o_got, _ = _run_paged( |
| 782 | inp, |
| 783 | cache, |
| 784 | idx, |
| 785 | warm, |
| 786 | use_fused=True, |
| 787 | checkpoint_state=cache, |
| 788 | checkpoint_offsets=torch.tensor([512], device="cuda", dtype=torch.int32), |
| 789 | checkpoint_state_indices=ckidx, |
| 790 | ) |
| 791 | torch.testing.assert_close(o_got, o_ref, rtol=5e-3, atol=5e-3) |
| 792 | torch.testing.assert_close(cache[idx.long()], ht_ref, rtol=5e-3, atol=5e-3) |
| 793 | torch.testing.assert_close(cache[ckidx.long()], ck_ref, rtol=5e-3, atol=5e-3) |
| 794 | |
| 795 | |
| 796 | def test_paged_state_rejects_a_dense_initial_state() -> None: |
| 797 | """The two ways of passing the state are mutually exclusive.""" |
| 798 | _requires_kernel() |
| 799 | inp = _inputs([256], seed=15) |
| 800 | cache, idx, warm = _paged_cache(inp, [True], [1], slots=4) |
| 801 | with pytest.raises(ValueError, match="replaces initial_state"): |
| 802 | _run_paged( |
| 803 | inp, |
| 804 | cache, |
| 805 | idx, |
| 806 | warm, |
| 807 | use_fused=True, |
| 808 | initial_state=inp["h0"].clone(), |
| 809 | ) |
| 810 | |
| 811 | |
| 812 | @pytest.mark.parametrize("use_fused", [False, True]) |
| 813 | def test_padded_cache_stride_is_honoured(use_fused: bool) -> None: |
| 814 | """A mamba page may pad stride(0); only each state row is dense. |
| 815 | |
| 816 | The real `recurrent_state` is not a contiguous tensor, so a kernel that |
| 817 | assumes a packed ``H * V * K`` row stride addresses the wrong rows. The |
| 818 | other cases in this file all build a contiguous cache, which cannot |
| 819 | exercise that. |
| 820 | """ |
| 821 | _requires_kernel() |
| 822 | seqlens = [1024, 512] |
| 823 | inp = _inputs(seqlens, seed=21) |
| 824 | n = len(seqlens) |
| 825 | row = NUM_HEADS * HEAD_DIM * HEAD_DIM |
| 826 | slots, pad = 16, 512 |
| 827 | raw = torch.full((slots, row + pad), -7.0, device="cuda", dtype=torch.float32) |
| 828 | cache = raw.as_strided( |
| 829 | (slots, NUM_HEADS, HEAD_DIM, HEAD_DIM), |
| 830 | (row + pad, HEAD_DIM * HEAD_DIM, HEAD_DIM, 1), |
| 831 | ) |
| 832 | assert not cache.is_contiguous() and cache[0].is_contiguous() |
| 833 | |
| 834 | rows = [11, 3] |
| 835 | idx = torch.tensor(rows, device="cuda", dtype=torch.int32) |
| 836 | warm = torch.ones(n, dtype=torch.bool, device="cuda") |
| 837 | for i, r in enumerate(rows): |
| 838 | cache[r] = inp["h0"][i] |
| 839 | before = raw.clone() |
| 840 | |
| 841 | o_ref, ht_ref = _run(inp, use_fused=use_fused) |
| 842 | o_got, ht_got = _run_paged(inp, cache, idx, warm, use_fused=use_fused) |
| 843 | |
| 844 | assert ht_got is None |
| 845 | torch.testing.assert_close(o_got, o_ref, rtol=5e-3, atol=5e-3) |
| 846 | torch.testing.assert_close(cache[idx.long()], ht_ref, rtol=5e-3, atol=5e-3) |
| 847 | |
| 848 | # The padding between rows, and every unowned row, must be untouched. |
| 849 | untouched = torch.ones(slots, dtype=torch.bool, device="cuda") |
| 850 | untouched[idx.long()] = False |
| 851 | assert torch.equal(raw[untouched], before[untouched]) |
| 852 | assert torch.equal(raw[idx.long(), row:], before[idx.long(), row:]), ( |
| 853 | "the kernel wrote into the inter-row padding" |
| 854 | ) |
| 855 | |
| 856 | |
| 857 | def test_padded_checkpoint_stride_is_honoured() -> None: |
| 858 | """The checkpoint destination is the paged cache, which pads stride(0). |
| 859 | |
| 860 | Every other checkpoint case here builds a contiguous buffer, so none of |
| 861 | them would catch a kernel that assumes a packed ``H * V * K`` row stride |
| 862 | and writes the snapshot into a neighbouring block. |
| 863 | """ |
| 864 | _requires_kernel() |
| 865 | seqlens = [1024] |
| 866 | inp = _inputs(seqlens, seed=31) |
| 867 | row = NUM_HEADS * HEAD_DIM * HEAD_DIM |
| 868 | slots, pad = 16, 512 |
| 869 | raw = torch.full((slots, row + pad), -7.0, device="cuda", dtype=torch.float32) |
| 870 | cache = raw.as_strided( |
| 871 | (slots, NUM_HEADS, HEAD_DIM, HEAD_DIM), |
| 872 | (row + pad, HEAD_DIM * HEAD_DIM, HEAD_DIM, 1), |
| 873 | ) |
| 874 | assert not cache.is_contiguous() |
| 875 | |
| 876 | state_rows = [11] |
| 877 | idx = torch.tensor(state_rows, device="cuda", dtype=torch.int32) |
| 878 | warm = torch.ones(1, dtype=torch.bool, device="cuda") |
| 879 | cache[state_rows[0]] = inp["h0"][0] |
| 880 | ck_rows = [4] |
| 881 | before = raw.clone() |
| 882 | |
| 883 | prefix = _prefix_inputs(inp, 0, 512) |
| 884 | _, ck_ref = _run(prefix, use_fused=True) |
| 885 | o_ref, ht_ref = _run(inp, use_fused=True) |
| 886 | |
| 887 | o_got, _ = _run_paged( |
| 888 | inp, |
| 889 | cache, |
| 890 | idx, |
| 891 | warm, |
| 892 | use_fused=True, |
| 893 | checkpoint_state=cache, |
| 894 | checkpoint_offsets=torch.tensor([512], device="cuda", dtype=torch.int32), |
| 895 | checkpoint_state_indices=torch.tensor( |
| 896 | ck_rows, device="cuda", dtype=torch.int32 |
| 897 | ), |
| 898 | ) |
| 899 | |
| 900 | torch.testing.assert_close(o_got, o_ref, rtol=5e-3, atol=5e-3) |
| 901 | torch.testing.assert_close(cache[idx.long()], ht_ref, rtol=5e-3, atol=5e-3) |
| 902 | torch.testing.assert_close(cache[ck_rows[0]], ck_ref[0], rtol=5e-3, atol=5e-3) |
| 903 | |
| 904 | touched = torch.zeros(slots, dtype=torch.bool, device="cuda") |
| 905 | touched[state_rows + ck_rows] = True |
| 906 | assert torch.equal(raw[~touched], before[~touched]) |
| 907 | assert torch.equal(raw[touched, row:], before[touched, row:]), ( |
| 908 | "the kernel wrote into the inter-row padding" |
| 909 | ) |