@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 3e2a7e74a5

[Feat][Model] Enable KDA prefill checkpoints for GLM-5.3-Flash (#56960)

Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>

Chauncey committed Sep 29, 2026 at 17:14 UTC 3e2a7e74a556623ddee3a327299d7bf04055501e

6 files changed +307 -11

tests/models/glm5next/test_kda_recurrent.py

+167

@@ -193,3 +193,170 @@ def test_fused_recurrent_kda_rejects_unaddressable_layouts():

193 broken["cu_seqlens"] = None if q.shape[0] > 1 else inputs["cu_seqlens"]

194 with pytest.raises(AssertionError, match=r"torch.Size"):

195 run_kernel(broken, state)

196 +

197 +

198 +@pytest.mark.parametrize("offset", [16, 144, 256])

199 +@pytest.mark.parametrize("dim_first", [False, True])

200 +@pytest.mark.parametrize("num_spec", [0, 3])

201 +@torch.inference_mode()

202 +def test_prefill_checkpoint_resumes_suffix(monkeypatch, dim_first, num_spec, offset):

203 + """Restoring both cached states must reproduce an uninterrupted prefill."""

204 + from types import SimpleNamespace

205 +

206 + from vllm.model_executor.layers.mamba.checkpoint import (

207 + MambaPrefillCheckpointMetadata,

208 + )

209 + from vllm.model_executor.layers.mamba.kda_checkpoint import (

210 + FlashKDAPrefillCheckpointExporter,

211 + )

212 + from vllm.models.glm5next.common import kda

213 + from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata

214 + from vllm.v1.attention.backends.utils import compute_causal_conv1d_metadata

215 + from vllm.v1.worker.workspace import WorkspaceManager

216 +

217 + if kda._resolve_kda_prefill_backend("auto", D, torch.bfloat16, -5) != "flashkda":

218 + pytest.skip("requires FlashKDA")

219 + import vllm._flashkda_C # noqa: F401

220 +

221 + torch.manual_seed(42)

222 + device = torch.device("cuda")

223 + length, heads = 273, 2

224 + width, state_len = 3 * heads * D, 3 + num_spec

225 + layer = object.__new__(kda.Glm5NextLinearAttention)

226 + torch.nn.Module.__init__(layer)

227 + layer.prefix = "layer.0"

228 + layer.local_num_heads = heads

229 + layer.head_dim = D

230 + layer.conv_size = 4

231 + layer.local_projection_size = heads * D

232 + layer.kda_safe_gate = True

233 + layer.kda_lower_bound = -5.0

234 + layer.kda_prefill_backend = "flashkda"

235 + layer._conv_state_dim_first = dim_first

236 + layer._merged_conv_weight = torch.randn(width, 4, device=device) * 0.1

237 + layer.q_conv1d = SimpleNamespace(bias=None)

238 + layer.A_log = torch.randn(1, 1, heads, 1, device=device) * 0.1

239 + layer.dt_bias = torch.randn(heads * D, device=device) * 0.1

240 + state_shape = (2, heads, D, D)

241 + layer._flashkda_buffer_specs = (

242 + (state_shape, torch.float32),

243 + (state_shape, torch.float32),

244 + (

245 + (torch.ops._flashkda_C.get_workspace_size(2 * length, heads, 2),),

246 + torch.uint8,

247 + ),

248 + ((1, 2 * length, heads, D), torch.bfloat16),

249 + )

250 + layer._checkpoint_exporter = FlashKDAPrefillCheckpointExporter(state_len=3)

251 + workspace = WorkspaceManager(device)

252 + monkeypatch.setattr(kda, "current_workspace_manager", lambda: workspace)

253 + conv_shape = (width, state_len) if dim_first else (state_len, width)

254 + num_slots = 5 + (num_spec + 1 if num_spec else 0)

255 + conv_storage = torch.zeros(

256 + num_slots, width * state_len + 8, device=device, dtype=torch.bfloat16

257 + )

258 + conv = conv_storage[:, : width * state_len].view(num_slots, *conv_shape)

259 + recurrent_storage = torch.zeros(

260 + num_slots, heads * D * D + 8, device=device, dtype=torch.float32

261 + )

262 + recurrent = recurrent_storage[:, : heads * D * D].view(num_slots, heads, D, D)

263 + layer.kv_cache = (conv, recurrent)

264 + projected = torch.randn(

265 + 2 * length, width + heads, device=device, dtype=torch.bfloat16

266 + )

267 + raw_qkv, beta = projected[:, :width], projected[:, width:].unsqueeze(0)

268 + raw_before = raw_qkv.clone()

269 + gate = torch.randn(1, 2 * length, heads, D, device=device, dtype=torch.bfloat16)

270 +

271 + def run(qkv, g, b, slots, has_initial, checkpoint=None):

272 + n = qkv.shape[0] // 2

273 + cu_cpu = torch.tensor([0, n, 2 * n], dtype=torch.int32)

274 + nums, batch_ptr, chunk_offsets = compute_causal_conv1d_metadata(

275 + cu_cpu, device=device

276 + )

277 + metadata = GDNAttentionMetadata(

278 + num_prefills=2,

279 + num_prefill_tokens=2 * n,

280 + num_decodes=0,

281 + num_decode_tokens=0,

282 + num_spec_decodes=0,

283 + num_spec_decode_tokens=0,

284 + num_actual_tokens=2 * n,

285 + has_initial_state=torch.tensor(has_initial, device=device),

286 + non_spec_query_start_loc=cu_cpu.to(device),

287 + non_spec_state_indices_tensor=torch.tensor(

288 + slots, dtype=torch.int32, device=device

289 + ),

290 + checkpoint=checkpoint,

291 + nums_dict=nums,

292 + batch_ptr=batch_ptr,

293 + token_chunk_offset_ptr=chunk_offsets,

294 + )

295 + monkeypatch.setattr(

296 + kda,

297 + "get_forward_context",

298 + lambda: SimpleNamespace(attn_metadata={layer.prefix: metadata}),

299 + )

300 + spec_tokens = num_spec + 1 if num_spec else 0

301 + if spec_tokens:

302 + qkv = torch.cat([qkv[:spec_tokens], qkv])

303 + g = torch.cat([g[:, :spec_tokens], g], dim=1)

304 + b = torch.cat([b[:, :spec_tokens], b], dim=1)

305 + metadata.num_spec_decodes = 1

306 + metadata.num_spec_decode_tokens = spec_tokens

307 + metadata.num_actual_tokens += spec_tokens

308 + metadata.spec_sequence_masks = torch.tensor(

309 + [True, False, False], device=device

310 + )

311 + metadata.spec_query_start_loc = torch.tensor(

312 + [0, spec_tokens], dtype=torch.int32, device=device

313 + )

314 + metadata.spec_state_indices_tensor = torch.arange(

315 + 5, num_slots, dtype=torch.int32, device=device

316 + ).unsqueeze(0)

317 + metadata.spec_token_indx = torch.arange(spec_tokens, device=device)

318 + metadata.non_spec_token_indx = torch.arange(

319 + spec_tokens, qkv.shape[0], device=device

320 + )

321 + metadata.num_accepted_tokens = torch.ones(

322 + 1, dtype=torch.int32, device=device

323 + )

324 + out = torch.empty(

325 + 1, qkv.shape[0], heads, D, device=device, dtype=torch.bfloat16

326 + )

327 + layer._forward(qkv, g, b, out)

328 + return out[:, spec_tokens:]

329 +

330 + checkpoint = MambaPrefillCheckpointMetadata(

331 + torch.tensor([offset, 0], dtype=torch.int32, device=device),

332 + torch.tensor([3, 0], dtype=torch.int32, device=device),

333 + )

334 + expected = run(raw_qkv, gate, beta, [1, 2], [False, False], checkpoint)

335 + expected_state = recurrent[1].clone()

336 + saved_conv, saved_recurrent = conv[3].clone(), recurrent[3].clone()

337 + torch.testing.assert_close(raw_qkv, raw_before)

338 + conv_view = conv if dim_first else conv.transpose(-1, -2)

339 + torch.testing.assert_close(conv_view[3, :, :3], raw_qkv[offset - 3 : offset].T)

340 + assert torch.count_nonzero(conv_view[3, :, 3:]) == 0

341 + assert torch.count_nonzero(conv[0]) == torch.count_nonzero(recurrent[0]) == 0

342 +

343 + run(

344 + torch.cat([raw_qkv[:offset], raw_qkv[length : length + offset]]),

345 + torch.cat([gate[:, :offset], gate[:, length : length + offset]], dim=1),

346 + torch.cat([beta[:, :offset], beta[:, length : length + offset]], dim=1),

347 + [1, 2],

348 + [False, False],

349 + )

350 + torch.testing.assert_close(saved_recurrent, recurrent[1], atol=0.003, rtol=0.03)

351 +

352 + suffix = slice(offset, length)

353 + qkv = torch.cat([raw_qkv[suffix], raw_qkv[length + offset :]])

354 + g = torch.cat([gate[:, suffix], gate[:, length + offset :]], dim=1)

355 + b = torch.cat([beta[:, suffix], beta[:, length + offset :]], dim=1)

356 + conv[4].copy_(saved_conv)

357 + recurrent[4].copy_(saved_recurrent)

358 + actual = run(qkv, g, b, [3, 4], [True, True])

359 + torch.testing.assert_close(

360 + actual[:, : length - offset], expected[:, suffix], atol=0.003, rtol=0.03

361 + )

362 + torch.testing.assert_close(recurrent[3], expected_state, atol=0.003, rtol=0.03)

tests/v1/attention/test_gdn_metadata_builder.py

+63

@@ -367,3 +367,66 @@ def test_cudagraph_capture_batch_stays_decode_only():

367 assert staged is not None

368 assert staged.data_ptr() == builder.non_spec_state_indices_tensor.data_ptr()

369 torch.testing.assert_close(staged, common_attn_metadata.block_table_tensor[:, 0])

370 +

371 +

372 +@pytest.mark.parametrize("num_spec", [0, 3])

373 +@pytest.mark.parametrize("prefix_match_unit, expected_offset", [(16, 96), (8, 96)])

374 +def test_checkpoint_metadata_preserves_non_spec_order(

375 + num_spec, prefix_match_unit, expected_offset

376 +):

377 + """Only eligible non-spec rows get a checkpoint in their reserved page."""

378 + from vllm.v1.attention.backends.utils import NULL_BLOCK_ID

379 +

380 + vllm_config = create_vllm_config(

381 + model_name="Qwen/Qwen3.5-0.8B",

382 + block_size=BLOCK_SIZE,

383 + )

384 + if num_spec > 0:

385 + vllm_config.speculative_config = SpeculativeConfig(

386 + method="ngram",

387 + num_speculative_tokens=num_spec,

388 + )

389 + vllm_config.cache_config.mamba_cache_mode = "align"

390 + vllm_config.cache_config.prefix_match_unit = prefix_match_unit

391 + mamba_spec = MambaSpec(

392 + block_size=64,

393 + shapes=((16, 64),),

394 + dtypes=(torch.float16,),

395 + mamba_cache_mode="align",

396 + num_speculative_blocks=num_spec,

397 + num_prefill_checkpoint_blocks=1,

398 + prefill_checkpoint_alignment=16,

399 + )

400 + builder = GDNAttentionMetadataBuilder(

401 + kv_cache_spec=mamba_spec,

402 + layer_names=["layer.0"],

403 + vllm_config=vllm_config,

404 + device=DEVICE,

405 + )

406 + batch = BatchSpec(

407 + seq_lens=[65, 100, 100], query_lens=[3 if num_spec else 1, 100, 99]

408 + )

409 + common = create_common_attn_metadata(batch, 64, DEVICE, arange_block_indices=True)

410 + table = common.block_table_tensor + 1

411 + if num_spec:

412 + table = torch.cat([table, torch.zeros(3, num_spec, dtype=table.dtype)], dim=1)

413 + common = common.replace(block_table_tensor=table)

414 + kwargs = {}

415 + if num_spec:

416 + kwargs = dict(

417 + num_decode_draft_tokens_cpu=torch.tensor([2, -1, -1], dtype=torch.int32),

418 + num_accepted_tokens=torch.ones(3, dtype=torch.int32),

419 + )

420 +

421 + actual = builder.build(0, common, **kwargs)

422 + assert actual.checkpoint is not None

423 + offsets = [expected_offset, 0] if num_spec else [0, expected_offset, 0]

424 + slots = [common.block_table_tensor[1, 0].item(), NULL_BLOCK_ID]

425 + if not num_spec:

426 + slots.insert(0, NULL_BLOCK_ID)

427 + torch.testing.assert_close(

428 + actual.checkpoint.checkpoint_offsets, torch.tensor(offsets, dtype=torch.int32)

429 + )

430 + torch.testing.assert_close(

431 + actual.checkpoint.state_indices, torch.tensor(slots, dtype=torch.int32)

432 + )

vllm/model_executor/layers/mamba/checkpoint.py

@@ -76,6 +76,8 @@ class MambaPrefillCheckpointBuilder:

76 m: CommonAttentionMetadata,

77 request_rows: list[int],

78 ) -> MambaPrefillCheckpointMetadata | None:

79 + if self.vllm_config.cache_config.mamba_cache_mode != "align":

80 + return None

81 if self.kv_cache_spec.num_prefill_checkpoint_blocks == 0:

82 return None

83 assert m.seq_lens_cpu_upper_bound is not None

vllm/models/glm5next/common/kda.py

+59 -3

@@ -2,6 +2,8 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """GLM-5.3-Flash KDA layer with separate convolutions and a bounded safe gate."""

5 +from dataclasses import replace

6 +

7 import torch

8 from torch import nn

@@ -14,7 +16,14 @@ from vllm.model_executor.layers.linear import (

16 MergedColumnParallelLinear,

17 RowParallelLinear,

18 )

19 +from vllm.model_executor.layers.mamba.checkpoint import (

20 + MambaPrefillCheckpointMetadata,

21 +)

22 from vllm.model_executor.layers.mamba.gdn.base import GatedDeltaNetAttention

23 +from vllm.model_executor.layers.mamba.kda_checkpoint import (

24 + FlashKDAPrefillCheckpointExporter,

25 + kda_prefill_checkpoint_alignment,

26 +)

27 from vllm.model_executor.layers.mamba.mamba_utils import (

28 MambaStateDtypeCalculator,

29 MambaStateShapeCalculator,

@@ -37,6 +46,7 @@ from vllm.platforms import current_platform

46 from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated

47 from vllm.transformers_utils.configs.glm5_next import Glm5NextConfig

48 from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata

49 +from vllm.v1.kv_cache_interface import MambaSpec

50 from vllm.v1.worker.workspace import current_workspace_manager

52 if current_platform.is_rocm():

@@ -155,6 +165,16 @@ class Glm5NextLinearAttention(GatedDeltaNetAttention):

165 num_heads: int

166 conv_size: int

167

168 + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> MambaSpec:

169 + spec = super().get_kv_cache_spec(vllm_config)

170 + assert isinstance(spec, MambaSpec)

171 + alignment = kda_prefill_checkpoint_alignment(self.kda_prefill_backend)

172 + return replace(

173 + spec,

174 + num_prefill_checkpoint_blocks=int(alignment is not None),

175 + prefill_checkpoint_alignment=alignment,

176 + )

177 +

178 def get_state_dtype(

179 self,

180 ) -> tuple[torch.dtype, torch.dtype]:

@@ -330,6 +350,7 @@ class Glm5NextLinearAttention(GatedDeltaNetAttention):

350 self._flashkda_buffer_specs: (

351 tuple[tuple[tuple[int, ...], torch.dtype], ...] | None

352 ) = None

353 + self._checkpoint_exporter: FlashKDAPrefillCheckpointExporter | None = None

354 if self.kda_prefill_backend == "flashkda":

355 import vllm._flashkda_C # noqa: F401

356

@@ -339,6 +360,10 @@ class Glm5NextLinearAttention(GatedDeltaNetAttention):

360 max_tokens, self.local_num_heads, max_seqs

361 )

362 self._flashkda_buffer_specs = (

363 + (

364 + (max_seqs, self.local_num_heads, self.head_dim, self.head_dim),

365 + self.get_state_dtype()[1],

366 + ),

367 (

368 (max_seqs, self.local_num_heads, self.head_dim, self.head_dim),

369 self.get_state_dtype()[1],

@@ -351,6 +376,9 @@ class Glm5NextLinearAttention(GatedDeltaNetAttention):

376 vllm_config.model_config.dtype,

377 ),

378 )

379 + self._checkpoint_exporter = FlashKDAPrefillCheckpointExporter(

380 + state_len=self.conv_size - 1

381 + )

382

383 def _flashkda_prefill(

384 self,

@@ -362,6 +390,10 @@ class Glm5NextLinearAttention(GatedDeltaNetAttention):

390 initial_state: torch.Tensor,

391 cu_seqlens: torch.Tensor,

392 out: torch.Tensor | None,

393 + checkpoint: MambaPrefillCheckpointMetadata | None = None,

394 + raw_qkv: torch.Tensor | None = None,

395 + conv_state: torch.Tensor | None = None,

396 + recurrent_state: torch.Tensor | None = None,

397 ) -> tuple[torch.Tensor, torch.Tensor]:

398 """Fused KDA chunked prefill (FlashKDA). Takes the raw gate logits ``g``

399 and raw ``beta`` logits, l2-normalizes q/k in-kernel and applies the

@@ -370,10 +402,15 @@ class Glm5NextLinearAttention(GatedDeltaNetAttention):

402 attention output into ``out`` (a workspace buffer when ``None``) and

403 returns ``(out, final_state)``."""

404 assert self._flashkda_buffer_specs is not None

373 - final_state, workspace, workspace_out = (

405 + final_state, checkpoint_state, workspace, workspace_out = (

406 current_workspace_manager().get_simultaneous(*self._flashkda_buffer_specs)

407 )

408 final_state = final_state[: initial_state.shape[0]]

409 + checkpoint_state = (

410 + checkpoint_state[: initial_state.shape[0]]

411 + if checkpoint is not None

412 + else None

413 + )

414 if out is None:

415 out = workspace_out[:, : q.shape[1]]

416 # FlashKDA hardcodes dense q/k/v/g strides; beta may be row-strided.

@@ -392,9 +429,23 @@ class Glm5NextLinearAttention(GatedDeltaNetAttention):

429 initial_state.contiguous(),

430 final_state,

431 cu_seqlens.contiguous(),

395 - None,

396 - None,

432 + checkpoint_state,

433 + checkpoint.checkpoint_offsets.contiguous()

434 + if checkpoint is not None

435 + else None,

436 )

437 + if checkpoint is not None:

438 + assert self._checkpoint_exporter is not None

439 + assert raw_qkv is not None and conv_state is not None

440 + assert recurrent_state is not None

441 + self._checkpoint_exporter.export(

442 + checkpoint,

443 + raw_qkv=raw_qkv,

444 + conv_state=conv_state,

445 + recurrent_checkpoint=checkpoint_state,

446 + recurrent_state=recurrent_state,

447 + cu_seqlens=cu_seqlens,

448 + )

449 return out, final_state

450

451 def forward(

@@ -565,6 +616,7 @@ class Glm5NextLinearAttention(GatedDeltaNetAttention):

616 q_spec, k_spec, v_spec = qkv_spec.split(self.local_projection_size, dim=-1)

617

618 # --- causal conv1d: non-spec path (prefill or plain decode) ---

619 + raw_qkv_ns = qkv_ns

620 q_ns = k_ns = v_ns = None

621 if attn_metadata_narrowed.num_prefills > 0:

622 assert qkv_ns is not None

@@ -662,6 +714,10 @@ class Glm5NextLinearAttention(GatedDeltaNetAttention):

714 initial_state=initial_state,

715 cu_seqlens=non_spec_query_start_loc,

716 out=ns_out,

717 + checkpoint=attn_metadata_narrowed.checkpoint,

718 + raw_qkv=raw_qkv_ns,

719 + conv_state=conv_state,

720 + recurrent_state=recurrent_state,

721 )

722 else:

723 (

vllm/models/kimi_k3/nvidia/kda_metadata.py

-8

@@ -18,10 +18,6 @@ from typing import TYPE_CHECKING

18 import torch

20 from vllm.config import VllmConfig

21 -from vllm.model_executor.layers.mamba.checkpoint import (

22 - MambaPrefillCheckpointBuilder,

23 - MambaPrefillCheckpointMetadata,

24 -)

21 from vllm.platforms import current_platform

22 from vllm.triton_utils import tl, triton

23 from vllm.utils.torch_utils import async_tensor_h2d

@@ -271,7 +267,6 @@ class KimiK3KDAMetadata(GDNAttentionMetadata, RecoverSSMMetadata):

267 recoverssm_context: "KDARecoverSSMCommitContext | None" = field(

268 default=None, repr=False, compare=False

269 )

274 - checkpoint: MambaPrefillCheckpointMetadata | None = None

270

271 def commit_recoverssm_state(

272 self, num_accepted_tokens: torch.Tensor

@@ -319,9 +314,6 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):

314 device: torch.device,

315 ) -> None:

316 super().__init__(kv_cache_spec, layer_names, vllm_config, device)

322 - self.checkpoint_builder = MambaPrefillCheckpointBuilder(

323 - vllm_config, kv_cache_spec

324 - )

317 additional_config = vllm_config.additional_config

318 self.use_flashinfer_prefill = (

319 isinstance(additional_config, dict)

vllm/v1/attention/backends/gdn_attn.py

+16

@@ -8,6 +8,10 @@ from typing import Literal

8 import torch

10 from vllm.config import VllmConfig

11 +from vllm.model_executor.layers.mamba.checkpoint import (

12 + MambaPrefillCheckpointBuilder,

13 + MambaPrefillCheckpointMetadata,

14 +)

15 from vllm.utils.torch_utils import async_tensor_h2d

16 from vllm.v1.attention.backend import (

17 AttentionBackend,

@@ -48,6 +52,7 @@ class GDNAttentionMetadata:

52 num_spec_decode_tokens: int

53 num_actual_tokens: int

55 + checkpoint: MambaPrefillCheckpointMetadata | None = None

56 has_initial_state: torch.Tensor | None = None

58 spec_query_start_loc: torch.Tensor | None = None # shape: [num_spec_decodes + 1,]

@@ -97,6 +102,9 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]

102 super().__init__(kv_cache_spec, layer_names, vllm_config, device)

103 self.compilation_config = vllm_config.compilation_config

104 self.speculative_config = vllm_config.speculative_config

105 + self.checkpoint_builder = MambaPrefillCheckpointBuilder(

106 + vllm_config, kv_cache_spec

107 + )

108 from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (

109 _resolve_gdn_prefill_backend,

110 )

@@ -466,6 +474,13 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]

474 else:

475 has_initial_state = None

476

477 + checkpoint: MambaPrefillCheckpointMetadata | None = None

478 + if num_prefills > 0:

479 + request_rows = list(range(m.num_reqs))

480 + if spec_sequence_masks_cpu is not None:

481 + request_rows = (~spec_sequence_masks_cpu).nonzero().flatten().tolist()

482 + checkpoint = self.checkpoint_builder.build(m, request_rows)

483 +

484 # Function code counted on either presency non-spec decode or spec decode,

485 # but not both.

486 assert not (num_decodes > 0 and num_spec_decodes > 0), (

@@ -543,6 +558,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]

558 num_spec_decodes=num_spec_decodes,

559 num_spec_decode_tokens=num_spec_decode_tokens,

560 num_actual_tokens=m.num_actual_tokens,

561 + checkpoint=checkpoint,

562 has_initial_state=has_initial_state,

563 chunk_indices=chunk_indices,

564 chunk_offsets=chunk_offsets,