@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 1417022c3d

[ROCm][Kimi-K3] Optimize low-concurrency speculative KDA (#58045)

Signed-off-by: jiacao-amd <jiahui.cao@amd.com> Co-authored-by: OpenAI Codex <noreply@openai.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>

jiacao-amd committed Sep 25, 2026 at 10:33 UTC 1417022c3dbb633f9d0c2ac1db2171445bb91b08

6 files changed +201 -27

tests/models/kimi_k3/test_kda.py

+56 -9

@@ -22,6 +22,9 @@ from vllm.model_executor.layers.mamba.ops.causal_conv1d import causal_conv1d_upd

22 from vllm.model_executor.layers.mamba.ops.gather_initial_states import (

23 gather_initial_states,

24 )

25 +from vllm.models.kimi_k3.amd.ops.third_party.kda import (

26 + fused_recurrent_kda as fused_recurrent_kda_amd,

27 +)

28 from vllm.models.kimi_k3.amd.ops.third_party.kda import (

29 fused_recurrent_kda_packed_decode as fused_recurrent_kda_packed_decode_amd,

30 )

@@ -44,10 +47,12 @@ from vllm.models.kimi_k3.nvidia.ops.third_party.kda import (

47 chunk_kda,

48 chunk_kda_with_fused_gate,

49 fused_kda_gate,

47 - fused_recurrent_kda,

50 fused_recurrent_kda_fwd,

51 fused_recurrent_kda_packed_decode,

52 )

53 +from vllm.models.kimi_k3.nvidia.ops.third_party.kda import (

54 + fused_recurrent_kda as fused_recurrent_kda_nvidia,

55 +)

56 from vllm.platforms import current_platform

57 from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd

58 from vllm.utils.flashinfer import flashinfer_fused_kda_decode

@@ -66,6 +71,10 @@ PACKED_DECODE_IMPLS = {

71 "nvidia": fused_recurrent_kda_packed_decode,

72 "amd": fused_recurrent_kda_packed_decode_amd,

73 }

74 +SPEC_DECODE_IMPLS = {

75 + "nvidia": fused_recurrent_kda_nvidia,

76 + "amd": fused_recurrent_kda_amd,

77 +}

80 def test_kda_warmup_skips_missing_metadata(monkeypatch):

@@ -468,17 +477,51 @@ def test_packed_kda_decode_correctness(

477

478

479 @pytest.mark.parametrize(

471 - ("H", "fuse_gate"),

472 - [(12, True), (12, False), (12, None), (96, None)],

480 + ("impl", "H", "fuse_gate", "num_seqs", "query_len"),

481 + [

482 + *[

483 + pytest.param(

484 + impl,

485 + H,

486 + fuse_gate,

487 + 3,

488 + 3,

489 + id=f"{impl}-H{H}-fuse-{fuse_gate}",

490 + )

491 + for impl in SPEC_DECODE_IMPLS

492 + for H, fuse_gate in [

493 + (12, True),

494 + (12, False),

495 + (12, None),

496 + (96, None),

497 + ]

498 + ],

499 + pytest.param("amd", 12, True, 1, 1, id="amd-single-token"),

500 + pytest.param("amd", 12, True, 1, 8, id="amd-single-sequence-short"),

501 + pytest.param(

502 + "amd",

503 + 96,

504 + None,

505 + 1,

506 + 8,

507 + id="amd-single-sequence-many-heads",

508 + ),

509 + pytest.param("amd", 12, True, 2, 4, id="amd-uniform-two-sequences"),

510 + pytest.param("amd", 12, True, 4, 4, id="amd-uniform-four-sequences"),

511 + pytest.param("amd", 12, True, 8, 7, id="amd-uniform-eight-sequences"),

512 + ],

513 )

514 @pytest.mark.parametrize("lower_bound", [-5.0, None])

515 @torch.inference_mode()

516 def test_kda_spec_decode_correctness(

517 H: int,

518 fuse_gate: bool | None,

519 + num_seqs: int,

520 + query_len: int,

521 lower_bound: float | None,

522 + impl: str,

523 ):

481 - num_seqs, query_len, D = 3, 3, 128

524 + D = 128

525 T = num_seqs * query_len

526 torch.manual_seed(1234)

527

@@ -522,10 +565,8 @@ def test_kda_spec_decode_correctness(

565 dtype=torch.int32,

566 device=DEVICE,

567 ).view(num_seqs, query_len)

525 - num_accepted_tokens = torch.tensor(

526 - [1, 2, 3],

527 - dtype=torch.int32,

528 - device=DEVICE,

568 + num_accepted_tokens = (

569 + torch.arange(num_seqs, dtype=torch.int32, device=DEVICE) % query_len + 1

570 )

571 state_storage = 0.01 * torch.randn(

572 T + 1,

@@ -578,7 +619,12 @@ def test_kda_spec_decode_correctness(

619 expected = torch.cat(expected_outputs, dim=1)

620

621 actual_state = state.clone()

581 - actual, _ = fused_recurrent_kda(

622 + extra_args = (

623 + {"uniform_sequence_length": query_len}

624 + if impl == "amd" and num_seqs in (2, 4, 8)

625 + else {}

626 + )

627 + actual, _ = SPEC_DECODE_IMPLS[impl](

628 q=q,

629 k=k,

630 v=v,

@@ -593,6 +639,7 @@ def test_kda_spec_decode_correctness(

639 num_accepted_tokens=num_accepted_tokens,

640 out=output,

641 fuse_gate=fuse_gate,

642 + **extra_args,

643 )

644

645 assert actual.data_ptr() == output.data_ptr()

tests/models/kimi_k3/test_kda_metadata.py

+19

@@ -49,6 +49,7 @@ PRUNED_METADATA_FIELDS = {

49 "prefill_state_indices",

50 "prefill_has_initial_state",

51 "spec_sequence_masks",

52 + "uniform_spec_sequence_length",

53 "flashinfer_prefill_query_start_loc",

54 "flashinfer_prefill_seq_order",

55 }

@@ -409,6 +410,7 @@ def test_internal_checkpoint_metadata_skips_unaligned_offset():

410 "num_speculative_tokens",

411 "full_cuda_graph",

412 "is_prefilling",

413 + "expected_uniform_spec_sequence_length",

414 ),

415 [

416 pytest.param(

@@ -417,6 +419,7 @@ def test_internal_checkpoint_metadata_skips_unaligned_offset():

419 2,

420 False,

421 [False, False],

422 + 3,

423 id="pure-spec-decode",

424 ),

425 pytest.param(

@@ -425,6 +428,7 @@ def test_internal_checkpoint_metadata_skips_unaligned_offset():

428 2,

429 False,

430 [True, False, False],

431 + 3,

432 id="mixed-prefill-and-spec-decode",

433 ),

434 pytest.param(

@@ -433,6 +437,7 @@ def test_internal_checkpoint_metadata_skips_unaligned_offset():

437 0,

438 False,

439 [False, False],

440 + None,

441 id="regular-decode",

442 ),

443 pytest.param(

@@ -441,8 +446,18 @@ def test_internal_checkpoint_metadata_skips_unaligned_offset():

446 2,

447 False,

448 [False, False],

449 + None,

450 id="no-scheduled-draft-tokens",

451 ),

452 + pytest.param(

453 + BatchSpec(seq_lens=[50, 30], query_lens=[3, 2]),

454 + [2, 1],

455 + 2,

456 + False,

457 + [False, False],

458 + None,

459 + id="ragged-spec-decode",

460 + ),

461 ],

462 )

463 def test_kimi_k3_kda_metadata_matches_shared_gdn(

@@ -451,6 +466,7 @@ def test_kimi_k3_kda_metadata_matches_shared_gdn(

466 num_speculative_tokens: int,

467 full_cuda_graph: bool,

468 is_prefilling: list[bool],

469 + expected_uniform_spec_sequence_length: int | None,

470 ):

471 kwargs: dict[str, torch.Tensor] = {}

472 if num_decode_draft_tokens is not None:

@@ -483,6 +499,9 @@ def test_kimi_k3_kda_metadata_matches_shared_gdn(

499

500 assert isinstance(actual, KimiK3KDAMetadata)

501 _assert_matches_shared_gdn(reference, actual)

502 + assert (

503 + reference.uniform_spec_sequence_length == expected_uniform_spec_sequence_length

504 + )

505

506

507 def test_mixed_regular_and_spec_decode_uses_packed_decode_metadata():

vllm/models/kimi_k3/amd/kda.py

@@ -438,6 +438,14 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):

438 for x in mixed_qkv_spec.split(self.local_projection_size, dim=-1)

439 )

440 spec_cu_seqlens = spec_query_start_loc[: m.num_spec_decodes + 1]

441 + uniform_sequence_length = m.uniform_spec_sequence_length

442 + # Token padding can make the packed tensor larger than the

443 + # request-local metadata represented by this length.

444 + if (

445 + uniform_sequence_length is not None

446 + and q_spec.shape[1] != m.num_spec_decodes * uniform_sequence_length

447 + ):

448 + uniform_sequence_length = None

449 # Spec-only batches write directly into core_attn_out.

450 spec_out = (

451 core_attn_out[:, : q_spec.shape[1]]

@@ -457,6 +465,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):

465 cu_seqlens=spec_cu_seqlens,

466 ssm_state_indices=spec_state_indices_tensor,

467 num_accepted_tokens=num_accepted_tokens,

468 + uniform_sequence_length=uniform_sequence_length,

469 out=spec_out,

470 )

471

vllm/models/kimi_k3/amd/ops/third_party/kda/__init__.py

+4 -3

@@ -16,9 +16,10 @@

16 # - OOB-mask correctness fix: present (all tl.load use mask=..., other=0).

17 # Validated on gfx950: no core-dump, gsm8k 94.1%.

18 #

19 -# AMD-specific deltas vs the NVIDIA copy: NONE yet (byte-identical). Keep in sync

20 -# with the NVIDIA copy on FLA updates; any divergence should be an intentional,

21 -# documented gfx950-specific change (a #869-style AMD-only fix).

19 +# AMD-specific deltas vs the NVIDIA copy:

20 +# - fused recurrent gated speculative-decode launch geometry is tuned for

21 +# gfx950.

22 +# Keep the copies in sync on FLA updates outside documented platform tuning.

24 from .chunk import (

25 chunk_kda,

vllm/models/kimi_k3/amd/ops/third_party/kda/fused_recurrent.py

+103 -15

@@ -164,15 +164,24 @@ def fused_recurrent_kda_fwd_kernel(

164 APPLY_BETA_SIGMOID: tl.constexpr,

165 HAS_DT_BIAS: tl.constexpr,

166 USE_LOWER_BOUND: tl.constexpr,

167 + IS_SINGLE_SEQUENCE: tl.constexpr,

168 + SEQUENCE_LENGTH: tl.constexpr,

169 num_stages: tl.constexpr,

170 ):

171 pid = tl.program_id(0)

172 i_v = pid % tl.cdiv(V, BV)

173 i_nh = pid // tl.cdiv(V, BV)

172 - i_n, i_h = i_nh // H, i_nh % H

173 - bos = tl.load(cu_seqlens + i_n).to(tl.int64)

174 - eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64)

175 - sequence_length = eos - bos

174 + if IS_SINGLE_SEQUENCE:

175 + i_n, i_h = 0, i_nh

176 + else:

177 + i_n, i_h = i_nh // H, i_nh % H

178 + if SEQUENCE_LENGTH:

179 + bos = i_n * SEQUENCE_LENGTH

180 + sequence_length = SEQUENCE_LENGTH

181 + else:

182 + bos = tl.load(cu_seqlens + i_n).to(tl.int64)

183 + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64)

184 + sequence_length = eos - bos

185 if sequence_length == 0:

186 return

187

@@ -208,6 +217,14 @@ def fused_recurrent_kda_fwd_kernel(

217 p_v = v + bos * stride_qkv_token + i_h * V + o_v

218 p_g = g + bos * stride_g_token + i_h * K + o_k

219 p_beta = beta + bos * stride_beta_token + i_h

220 + if USE_GATE_IN_KERNEL:

221 + b_a = exp(tl.load(A_log + i_h).to(tl.float32))

222 + if HAS_DT_BIAS:

223 + b_bias = tl.load(

224 + dt_bias + i_h * K + o_k,

225 + mask=m_k,

226 + other=0.0,

227 + ).to(tl.float32)

228 for i_t in tl.range(0, sequence_length, num_stages=num_stages):

229 b_q = tl.load(p_q, mask=m_k, other=0.0, eviction_policy="evict_last").to(

230 tl.float32

@@ -219,8 +236,8 @@ def fused_recurrent_kda_fwd_kernel(

236 tl.float32

237 )

238 if USE_QK_L2NORM_IN_KERNEL:

222 - b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6)

223 - b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6)

239 + b_q *= tl.rsqrt(tl.sum(b_q * b_q) + 1e-6)

240 + b_k *= tl.rsqrt(tl.sum(b_k * b_k) + 1e-6)

241 b_q *= scale

242

243 b_gate = tl.load(

@@ -231,13 +248,7 @@ def fused_recurrent_kda_fwd_kernel(

248 ).to(tl.float32)

249 if USE_GATE_IN_KERNEL:

250 if HAS_DT_BIAS:

234 - b_bias = tl.load(

235 - dt_bias + i_h * K + o_k,

236 - mask=m_k,

237 - other=0.0,

238 - ).to(tl.float32)

251 b_gate += b_bias

240 - b_a = exp(tl.load(A_log + i_h).to(tl.float32))

252 if USE_LOWER_BOUND:

253 b_gate = lower_bound * tl.sigmoid(b_a * b_gate)

254 else:

@@ -288,6 +299,60 @@ def fused_recurrent_kda_fwd_kernel(

299 p_out += stride_out_token

300

301

302 +def _get_static_sequence_length(

303 + total_tokens: int,

304 + head_sequences: int,

305 + is_single_sequence: bool,

306 + uniform_sequence_length: int | None,

307 +) -> int:

308 + if is_single_sequence:

309 + return total_tokens

310 + # MI355X performance gate: 192 head-sequences is 16 sequences at H=12.

311 + # Static loops win through 16; dynamic loops win at 32 for short queries.

312 + if (

313 + uniform_sequence_length is not None

314 + and uniform_sequence_length > 1

315 + and head_sequences <= 192

316 + ):

317 + return uniform_sequence_length

318 + return 0

319 +

320 +

321 +def _select_kda_launch_config(

322 + use_gate_in_kernel: bool,

323 + is_spec_decoding: bool,

324 + head_sequences: int,

325 + static_sequence_length: int,

326 + has_multi_token_sequence: bool,

327 +) -> tuple[int, int, int]:

328 + if static_sequence_length:

329 + BV = 4 if head_sequences <= 24 else 8

330 + if static_sequence_length == 1:

331 + num_stages = 2

332 + elif static_sequence_length <= 3:

333 + num_stages = 8 if head_sequences <= 24 else 4

334 + elif static_sequence_length <= 6:

335 + num_stages = 6

336 + elif static_sequence_length <= 10:

337 + num_stages = 10

338 + else:

339 + num_stages = 12

340 + return BV, 1, num_stages

341 +

342 + if use_gate_in_kernel and is_spec_decoding:

343 + if head_sequences <= 24:

344 + return 4, 1, 4 if has_multi_token_sequence else 2

345 + if has_multi_token_sequence and head_sequences <= 48:

346 + return 8, 1, 4

347 + if has_multi_token_sequence:

348 + return 8, 1, 2

349 + return 32, 4, 2

350 +

351 + if use_gate_in_kernel:

352 + return 32, 4, 2

353 + return 8, 1, 2

354 +

355 +

356 def fused_recurrent_kda_fwd(

357 q: torch.Tensor,

358 k: torch.Tensor,

@@ -300,6 +365,7 @@ def fused_recurrent_kda_fwd(

365 cu_seqlens: torch.Tensor | None = None,

366 ssm_state_indices: torch.Tensor | None = None,

367 num_accepted_tokens: torch.Tensor | None = None,

368 + uniform_sequence_length: int | None = None,

369 use_qk_l2norm_in_kernel: bool = True,

370 A_log: torch.Tensor | None = None,

371 dt_bias: torch.Tensor | None = None,

@@ -336,6 +402,10 @@ def fused_recurrent_kda_fwd(

402 assert num_accepted_tokens is None

403 else:

404 assert ssm_state_indices.stride(1) == 1

405 + if uniform_sequence_length is not None:

406 + assert num_accepted_tokens is not None

407 + assert uniform_sequence_length > 0

408 + assert N * uniform_sequence_length == T

409 assert cu_seqlens.is_contiguous()

410 if use_gate_in_kernel:

411 assert A_log is not None and A_log.is_contiguous()

@@ -344,8 +414,22 @@ def fused_recurrent_kda_fwd(

414 if scale is None:

415 scale = K**-0.5

416

347 - BV = 32 if use_gate_in_kernel else 8

348 - num_warps = 4 if use_gate_in_kernel else 1

417 + is_spec_decoding = num_accepted_tokens is not None

418 + is_single_sequence = is_spec_decoding and N == 1

419 + head_sequences = H * N

420 + static_sequence_length = _get_static_sequence_length(

421 + total_tokens=T,

422 + head_sequences=head_sequences,

423 + is_single_sequence=is_single_sequence,

424 + uniform_sequence_length=uniform_sequence_length,

425 + )

426 + BV, num_warps, num_stages = _select_kda_launch_config(

427 + use_gate_in_kernel=use_gate_in_kernel,

428 + is_spec_decoding=is_spec_decoding,

429 + head_sequences=head_sequences,

430 + static_sequence_length=static_sequence_length,

431 + has_multi_token_sequence=T > N,

432 + )

433 grid = (cdiv(V, BV) * N * H,)

434 fused_recurrent_kda_fwd_kernel[grid](

435 q=q,

@@ -379,8 +463,10 @@ def fused_recurrent_kda_fwd(

463 USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,

464 USE_GATE_IN_KERNEL=use_gate_in_kernel,

465 APPLY_BETA_SIGMOID=use_beta_sigmoid_in_kernel,

466 + IS_SINGLE_SEQUENCE=is_single_sequence,

467 + SEQUENCE_LENGTH=static_sequence_length,

468 num_warps=num_warps,

383 - num_stages=2,

469 + num_stages=num_stages,

470 )

471 return out, initial_state

472

@@ -398,6 +484,7 @@ def fused_recurrent_kda(

484 cu_seqlens: torch.Tensor,

485 ssm_state_indices: torch.Tensor,

486 num_accepted_tokens: torch.Tensor | None = None,

487 + uniform_sequence_length: int | None = None,

488 out: torch.Tensor | None = None,

489 fuse_gate: bool | None = None,

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

@@ -433,6 +520,7 @@ def fused_recurrent_kda(

520 cu_seqlens=cu_seqlens,

521 ssm_state_indices=ssm_state_indices,

522 num_accepted_tokens=num_accepted_tokens,

523 + uniform_sequence_length=uniform_sequence_length,

524 use_qk_l2norm_in_kernel=True,

525 A_log=A_log if fuse_gate else None,

526 dt_bias=dt_bias if fuse_gate else None,

vllm/v1/attention/backends/gdn_attn.py

+10

@@ -64,6 +64,7 @@ class GDNAttentionMetadata:

64 non_spec_token_indx: torch.Tensor | None = None

66 num_accepted_tokens: torch.Tensor | None = None # shape: [batch,]

67 + uniform_spec_sequence_length: int | None = None # None for ragged batches

69 # Pre-computed FLA chunk metadata (avoids GPU->CPU sync in prepare_chunk_indices)

70 chunk_indices: torch.Tensor | None = None

@@ -226,6 +227,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]

227 self.vllm_config.cache_config.mamba_cache_mode,

228 )

229

230 + uniform_spec_sequence_length = None

231 spec_sequence_masks_cpu: torch.Tensor | None = None

232 if not self.use_spec_decode or num_decode_draft_tokens_cpu is None:

233 spec_sequence_masks = None

@@ -289,6 +291,13 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]

291 assert spec_sequence_masks_cpu is not None

292 non_spec_sequence_masks_cpu = ~spec_sequence_masks_cpu

293 query_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]

294 + spec_query_lens_cpu = query_lens_cpu[spec_sequence_masks_cpu]

295 + if spec_query_lens_cpu.numel() > 0:

296 + first_spec_sequence_length = int(spec_query_lens_cpu[0])

297 + if first_spec_sequence_length > 0 and bool(

298 + torch.all(spec_query_lens_cpu == first_spec_sequence_length)

299 + ):

300 + uniform_spec_sequence_length = first_spec_sequence_length

301

302 # Use CPU tensors to avoid CPU-GPU sync

303 non_spec_query_lens_cpu = query_lens_cpu[non_spec_sequence_masks_cpu]

@@ -540,6 +549,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]

549 spec_token_indx=spec_token_indx,

550 non_spec_token_indx=non_spec_token_indx,

551 num_accepted_tokens=num_accepted_tokens,

552 + uniform_spec_sequence_length=uniform_spec_sequence_length,

553 nums_dict=nums_dict,

554 batch_ptr=batch_ptr,

555 token_chunk_offset_ptr=token_chunk_offset_ptr,