@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 5f7a20b316

[nv] add pcp support in dsv3.2 (#52046)

Signed-off-by: Summer Yang <girasoleyang@gmail.com>

Summer Yang committed Aug 18, 2026 at 14:59 UTC 5f7a20b3162ef88d531ada03aa1174643bb97c11

6 files changed +376 -119

tests/kernels/test_fused_deepseek_v32_norm_rope.py

+90 -6

@@ -292,6 +292,81 @@ def test_fused_norm_rope_no_indexer(num_tokens: int):

292 assert (topk == 7).all(), "topk buffer should be untouched on shared layer"

293

294

295 +@pytest.mark.parametrize("has_indexer", [False, True])

296 +def test_fused_norm_rope_materializes_pcp_cache_inputs(has_indexer: bool):

297 + """PCP gets local normalized/rotated K rows without direct cache writes."""

298 + torch.manual_seed(6)

299 + dev = "cuda"

300 + num_tokens = 17

301 + max_pos = 8192

302 + pos = torch.arange(num_tokens, device=dev, dtype=torch.int64)

303 + q_c = torch.randn(num_tokens, Q_LORA, device=dev, dtype=torch.bfloat16)

304 + kv_c = torch.randn(num_tokens, KV_LORA, device=dev, dtype=torch.bfloat16)

305 + k_pe = torch.randn(num_tokens, ROPE_DIM, device=dev, dtype=torch.bfloat16)

306 + qw = torch.randn(Q_LORA, device=dev, dtype=torch.bfloat16)

307 + kvw = torch.randn(KV_LORA, device=dev, dtype=torch.bfloat16)

308 + ik = (

309 + torch.randn(num_tokens, INDEX_HEAD_DIM, device=dev, dtype=torch.bfloat16)

310 + if has_indexer

311 + else None

312 + )

313 + ikw = (

314 + torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32)

315 + if has_indexer

316 + else None

317 + )

318 + ikb = (

319 + torch.randn(INDEX_HEAD_DIM, device=dev, dtype=torch.float32)

320 + if has_indexer

321 + else None

322 + )

323 + mla_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)

324 + idx_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev) if has_indexer else None

325 + q_out = torch.empty_like(q_c)

326 + kv_out = torch.empty_like(kv_c)

327 + kpe_out = torch.empty_like(k_pe)

328 + ik_out = torch.empty_like(ik) if ik is not None else None

329 + topk = torch.full((num_tokens, 2048), 7, device=dev, dtype=torch.int32)

330 +

331 + actual_q = K.fused_norm_rope(

332 + pos,

333 + q_c,

334 + qw,

335 + EPS,

336 + kv_c,

337 + kvw,

338 + EPS,

339 + k_pe,

340 + mla_cos_sin,

341 + ik,

342 + ikw,

343 + ikb,

344 + EPS,

345 + idx_cos_sin,

346 + topk,

347 + has_indexer=has_indexer,

348 + index_rope_interleave=True,

349 + q_c_out=q_out,

350 + kv_c_out=kv_out,

351 + k_pe_out=kpe_out,

352 + index_k_out=ik_out,

353 + )

354 +

355 + assert actual_q.data_ptr() == q_out.data_ptr()

356 + assert_bf16(actual_q, rms_norm(q_c, qw), "PCP q norm")

357 + assert_bf16(kv_out, rms_norm(kv_c, kvw), "PCP kv norm")

358 + assert_bf16(

359 + kpe_out,

360 + rope(k_pe.float(), pos, mla_cos_sin, interleave=True),

361 + "PCP k_pe RoPE",

362 + )

363 + if has_indexer:

364 + assert ik is not None and ikw is not None and ikb is not None

365 + assert ik_out is not None and idx_cos_sin is not None

366 + ik_ref = rope(layer_norm(ik, ikw, ikb), pos, idx_cos_sin, interleave=True)

367 + assert_bf16(ik_out, ik_ref, "PCP indexer-K")

368 +

369 +

370 @pytest.mark.parametrize("num_tokens", [1, 4, 17, 512])

371 def test_fused_norm_rope_ds_mla(num_tokens: int):

372 """fp8_ds_mla MLA cache layout (FlashMLA sparse, bf16-query path; SM90/SM100).

@@ -405,19 +480,29 @@ def test_fused_norm_rope_supports_large_token_count():

480 # ── fused_q ──────────────────────────────────────────────────────────────────

481

482

408 -@pytest.mark.parametrize("num_tokens", [1, 4, 17, 512, 4096])

483 +@pytest.mark.parametrize(

484 + ("num_tokens", "num_q_heads"),

485 + [

486 + (1, NUM_HEADS),

487 + (4, NUM_HEADS),

488 + (17, NUM_HEADS),

489 + (512, NUM_HEADS),

490 + (4096, NUM_HEADS),

491 + (17, 64),

492 + ],

493 +)

494 @pytest.mark.parametrize("index_interleave", [True, False])

410 -def test_fused_q(num_tokens: int, index_interleave: bool):

495 +def test_fused_q(num_tokens: int, num_q_heads: int, index_interleave: bool):

496 torch.manual_seed(2)

497 dev = "cuda"

498 max_pos = 8192

499 pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos

500

501 q_pe = torch.randn(

417 - num_tokens, NUM_HEADS, ROPE_DIM, device=dev, dtype=torch.bfloat16

502 + num_tokens, num_q_heads, ROPE_DIM, device=dev, dtype=torch.bfloat16

503 )

504 ql_nope = torch.randn(

420 - num_tokens, NUM_HEADS, KV_LORA, device=dev, dtype=torch.bfloat16

505 + num_tokens, num_q_heads, KV_LORA, device=dev, dtype=torch.bfloat16

506 )

507 index_q = torch.randn(

508 num_tokens, INDEX_HEADS, INDEX_HEAD_DIM, device=dev, dtype=torch.bfloat16

@@ -449,7 +534,7 @@ def test_fused_q(num_tokens: int, index_interleave: bool):

534 mqa_nope_ref = (ql_nope.float() / s).to(FP8)

535 qpe_ref = rope(

536 q_pe.float(),

452 - pos.unsqueeze(-1).expand(num_tokens, NUM_HEADS),

537 + pos.unsqueeze(-1).expand(num_tokens, num_q_heads),

538 q_cos_sin,

539 interleave=True,

540 )

@@ -484,7 +569,6 @@ def test_fused_q_no_indexer(num_tokens: int):

569 )

570 q_scale = torch.tensor([0.5], device=dev, dtype=torch.float32)

571 q_cos_sin = make_cos_sin(max_pos, ROPE_DIM, dev)

487 -

572 _, _, mqa = K.fused_q(

573 pos,

574 q_pe,

tests/model_executor/layers/test_mla_short_prefill_indexer.py

+45

@@ -177,3 +177,48 @@ def test_short_prefill_updates_k_cache_before_scoring_decision(

177 # K cache is always updated before the scoring decision.

178 torch.testing.assert_close(observed["k"], k[: slot_mapping.numel()])

179 assert observed["slots"] is slot_mapping

180 +

181 +

182 +def test_skipped_k_cache_insert_accepts_no_k(

183 + monkeypatch: pytest.MonkeyPatch,

184 +) -> None:

185 + indexer_metadata = make_indexer_metadata(

186 + num_prefills=0,

187 + num_prefill_tokens=0,

188 + slot_mapping=torch.empty(0, dtype=torch.long),

189 + )

190 + monkeypatch.setattr(

191 + sparse_indexer,

192 + "get_forward_context",

193 + lambda: SimpleNamespace(

194 + attn_metadata={INDEXER_LAYER: indexer_metadata},

195 + cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE,

196 + ),

197 + )

198 + monkeypatch.setattr(

199 + sparse_indexer.current_platform, "fp8_dtype", lambda: torch.float16

200 + )

201 +

202 + topk_indices = torch.full((1, 2048), 17, dtype=torch.int32)

203 + result = sparse_indexer.sparse_attn_indexer(

204 + torch.empty(1, 1),

205 + INDEXER_LAYER,

206 + torch.empty(1),

207 + torch.empty(1, 1),

208 + None,

209 + None,

210 + torch.empty(1, 1),

211 + 128,

212 + "ue8m0",

213 + 2048,

214 + 4,

215 + 4096,

216 + 4096,

217 + topk_indices,

218 + True,

219 + False,

220 + "",

221 + )

222 +

223 + assert result is topk_indices

224 + assert torch.all(topk_indices == -1)

vllm/model_executor/layers/sparse_attn_indexer.py

+6 -6

@@ -299,7 +299,7 @@ def sparse_attn_indexer(

299 kv_cache: torch.Tensor,

300 q_quant: torch.Tensor,

301 q_scale: torch.Tensor | None,

302 - k: torch.Tensor,

302 + k: torch.Tensor | None,

303 weights: torch.Tensor,

304 quant_block_size: int,

305 scale_fmt: str | None,

@@ -695,7 +695,7 @@ def sparse_attn_indexer_fake(

695 kv_cache: torch.Tensor,

696 q_quant: torch.Tensor,

697 q_scale: torch.Tensor | None,

698 - k: torch.Tensor,

698 + k: torch.Tensor | None,

699 weights: torch.Tensor,

700 quant_block_size: int,

701 scale_fmt: str | None,

@@ -781,7 +781,7 @@ class SparseAttnIndexer(CustomOp):

781 self,

782 hidden_states: torch.Tensor,

783 q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor],

784 - k: torch.Tensor,

784 + k: torch.Tensor | None,

785 weights: torch.Tensor,

786 ):

787 if current_platform.is_cuda() or current_platform.is_xpu():

@@ -798,7 +798,7 @@ class SparseAttnIndexer(CustomOp):

798 self,

799 hidden_states: torch.Tensor,

800 q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor],

801 - k: torch.Tensor,

801 + k: torch.Tensor | None,

802 weights: torch.Tensor,

803 ):

804 # FP8 path: single tensor (per-token scale is folded into `weights`).

@@ -835,7 +835,7 @@ class SparseAttnIndexer(CustomOp):

835 self,

836 hidden_states: torch.Tensor,

837 q_fp8: torch.Tensor,

838 - k: torch.Tensor,

838 + k: torch.Tensor | None,

839 weights: torch.Tensor,

840 ):

841 return self.forward_cuda(hidden_states, q_fp8, k, weights)

@@ -844,7 +844,7 @@ class SparseAttnIndexer(CustomOp):

844 self,

845 hidden_states: torch.Tensor,

846 q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor],

847 - k: torch.Tensor,

847 + k: torch.Tensor | None,

848 weights: torch.Tensor,

849 ):

850 assert not self.use_fp4_cache, "AMD platform doesn't support fp4 cache yet"

vllm/models/deepseek_v32/attention.py

+103 -20

@@ -1,5 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3 +from typing import TYPE_CHECKING, cast

4 +

5 import torch

6 import torch.nn as nn

7 from transformers import DeepseekV2Config, DeepseekV3Config

@@ -7,8 +9,14 @@ from transformers import DeepseekV2Config, DeepseekV3Config

9 from vllm.compilation.breakable_cudagraph import eager_break_during_capture

10 from vllm.config import CacheConfig, VllmConfig

11 from vllm.distributed import get_tensor_model_parallel_world_size

12 +from vllm.distributed.parallel_state import get_tp_group

13 from vllm.forward_context import get_forward_context

14 from vllm.model_executor.layers.attention import MLAAttention

15 +from vllm.model_executor.layers.attention.attention import get_attention_context

16 +from vllm.model_executor.layers.attention.pcp import (

17 + finalize_mla_pcp_decode,

18 + maybe_gather_mla_latent_cache_inputs,

19 +)

20 from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm

21 from vllm.model_executor.layers.linear import (

22 ColumnParallelLinear,

@@ -35,6 +43,9 @@ from vllm.models.deepseek_v32.common.kernels import fused_norm_rope, fused_q

43 from vllm.platforms import current_platform

44 from vllm.utils.torch_utils import is_quantized_kv_cache

46 +if TYPE_CHECKING:

47 + from vllm.model_executor.layers.attention.mla_attention import MLACommonMetadata

48 +

50 class DeepseekV32Indexer(nn.Module):

51 indexer_cache_cls = DeepseekV32IndexerCache

@@ -363,6 +374,7 @@ class DeepseekV32Attention(MLAAttention):

374 attn_metadata = attn_metadata_raw[0].get(self.layer_name)

375 else:

376 attn_metadata = attn_metadata_raw

377 + attn_metadata = cast("MLACommonMetadata | None", attn_metadata)

378

379 slot_mapping = forward_context.slot_mapping

380 assert isinstance(slot_mapping, dict)

@@ -374,7 +386,8 @@ class DeepseekV32Attention(MLAAttention):

386 indexer_k_norm_bias = self.indexer.k_norm.bias

387 indexer_k_norm_eps = self.indexer.k_norm.eps

388 indexer_k_rope_cos_sin_cache = self.indexer_rope_emb.cos_sin_cache

377 - indexer_k_cache = self.indexer.k_cache.kv_cache

389 + indexer_k_cache = None if self.use_pcp else self.indexer.k_cache.kv_cache

390 + index_k_out = torch.empty_like(index_k) if self.use_pcp else None

391 indexer_softmax_scale = self.indexer.softmax_scale

392 indexer_n_head_scale = self.indexer.n_head**-0.5

393 else:

@@ -384,10 +397,11 @@ class DeepseekV32Attention(MLAAttention):

397 indexer_k_norm_eps = 1e-6

398 indexer_k_rope_cos_sin_cache = None

399 indexer_k_cache = None

400 + index_k_out = None

401 indexer_softmax_scale = 0.0

402 indexer_n_head_scale = 0.0

403

390 - if attn_metadata is None:

404 + if attn_metadata is None or self.use_pcp:

405 mla_kv_cache = None

406 mla_k_scale = None

407 indexer_k_cache = None

@@ -396,6 +410,8 @@ class DeepseekV32Attention(MLAAttention):

410 mla_kv_cache = self.kv_cache

411 mla_k_scale = self._k_scale

412

413 + kv_c_out = torch.empty_like(kv_c) if self.use_pcp else None

414 + k_pe_out = torch.empty_like(k_pe) if self.use_pcp else None

415 q_c = fused_norm_rope(

416 positions,

417 q_c,

@@ -419,12 +435,14 @@ class DeepseekV32Attention(MLAAttention):

435 mla_k_scale=mla_k_scale,

436 has_indexer=has_indexer,

437 index_rope_interleave=self._index_rope_interleave,

438 + kv_c_out=kv_c_out,

439 + k_pe_out=k_pe_out,

440 + index_k_out=index_k_out,

441 )

442

443 q = self.q_b_proj(q_c)[0].view(-1, self.num_local_heads, self.qk_head_dim)

444 q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)

426 - q_nope = q_nope.transpose(0, 1)

427 - ql_nope = torch.bmm(q_nope, self.W_UK_T).transpose(0, 1)

445 + ql_nope = torch.bmm(q_nope.transpose(0, 1), self.W_UK_T).transpose(0, 1)

446

447 if self.indexer is not None and not self.skip_topk:

448 index_q = self.indexer.wq_b(q_c)[0]

@@ -451,7 +469,10 @@ class DeepseekV32Attention(MLAAttention):

469 self._sparse_indexer_and_attn(

470 q_c,

471 index_q_fp8,

472 + index_k_out,

473 index_weights_out,

474 + kv_c_out,

475 + k_pe_out,

476 ql_nope,

477 mqa_q,

478 output,

@@ -463,19 +484,26 @@ class DeepseekV32Attention(MLAAttention):

484 self,

485 q_c: torch.Tensor,

486 index_q_fp8: torch.Tensor | None,

487 + index_k: torch.Tensor | None,

488 index_weights_out: torch.Tensor | None,

489 + kv_c: torch.Tensor | None,

490 + k_pe: torch.Tensor | None,

491 ql_nope: torch.Tensor,

492 mqa_q: torch.Tensor,

493 output: torch.Tensor,

494 ) -> None:

495 if self.indexer is not None and not self.skip_topk:

496 + assert index_q_fp8 is not None

497 + assert index_weights_out is not None

498 + if self.use_pcp:

499 + assert index_k is not None

500 sparse_attn_indexer(

501 q_c,

502 self.indexer.k_cache.prefix,

503 self.indexer.k_cache.kv_cache,

504 index_q_fp8,

477 - None, # q_scale folded into weights on the fp8 path

478 - None, # k unused when skip_k_cache_insert=True

505 + None,

506 + index_k,

507 index_weights_out,

508 self.indexer.quant_block_size,

509 self.indexer.scale_fmt,

@@ -484,28 +512,56 @@ class DeepseekV32Attention(MLAAttention):

512 self.indexer.max_model_len,

513 self.indexer.max_total_seq_len,

514 self.topk_indices_buffer,

487 - skip_k_cache_insert=True,

488 - use_pcp=False,

515 + skip_k_cache_insert=not self.use_pcp,

516 + use_pcp=self.use_pcp,

517 dense_mha_metadata_layer_name=self._dense_mha_metadata_layer_name,

490 - use_fp4_cache=False,

491 - # fused_norm_rope already cleared the topk buffer this forward.

518 + dcp_rank=(

519 + self.dcp_manager.group.rank_in_group

520 + if self.dcp_manager is not None

521 + else 0

522 + ),

523 + dcp_world_size=(

524 + self._vllm_config.parallel_config.decode_context_parallel_size

525 + ),

526 + cp_kv_cache_interleave_size=(

527 + self._vllm_config.parallel_config.cp_kv_cache_interleave_size

528 + ),

529 skip_topk_buffer_clear=True,

530 )

531

495 - attn_metadata_raw = get_forward_context().attn_metadata

496 - if isinstance(attn_metadata_raw, dict):

497 - attn_metadata = attn_metadata_raw.get(self.layer_name)

498 - elif isinstance(attn_metadata_raw, list):

499 - attn_metadata = attn_metadata_raw[0].get(self.layer_name)

500 - else:

501 - attn_metadata = attn_metadata_raw

502 -

532 + attn_metadata, _, kv_cache, layer_slot_mapping = get_attention_context(

533 + self.layer_name

534 + )

535 if attn_metadata is None:

536 output.zero_()

537 return

538 + attn_metadata = cast("MLACommonMetadata", attn_metadata)

539 +

540 + if self.use_pcp:

541 + assert kv_c is not None and k_pe is not None

542 + kv_for_cache, kpe_for_cache, cache_slot_mapping = (

543 + maybe_gather_mla_latent_cache_inputs(

544 + kv_c,

545 + k_pe.unsqueeze(1),

546 + layer_slot_mapping,

547 + attn_metadata.num_decode_tokens,

548 + True,

549 + )

550 + )

551 + self.impl.do_kv_cache_update( # type: ignore[attr-defined]

552 + kv_for_cache,

553 + kpe_for_cache,

554 + kv_cache,

555 + cache_slot_mapping,

556 + self.kv_cache_dtype,

557 + self._k_scale,

558 + )

559

560 num_actual = attn_metadata.num_actual_tokens # type: ignore[attr-defined]

508 - kv_cache = self.kv_cache

561 + if num_actual == 0:

562 + output.zero_()

563 + return

564 +

565 if self._fp8_kv_needs_view:

566 kv_cache = kv_cache.view(torch.float8_e4m3fn)

567 if self._fp8_query:

@@ -515,10 +571,35 @@ class DeepseekV32Attention(MLAAttention):

571 ]

572 else:

573 mqa_q_arg = (ql_nope[:num_actual], mqa_q[:num_actual])

518 - attn_out, _ = self.impl.forward_mqa( # type: ignore[attr-defined]

574 +

575 + if self.use_pcp and self.impl.dcp_world_size > self.impl.pcp_world_size:

576 + if isinstance(mqa_q_arg, tuple):

577 + mqa_q_arg = torch.cat(mqa_q_arg, dim=-1)

578 + mqa_q_arg = get_tp_group().all_gather(mqa_q_arg, dim=1)

579 + attn_out, lse = self.impl.forward_mqa( # type: ignore[attr-defined]

580 mqa_q_arg, kv_cache, attn_metadata, self

581 )

582

583 + if self.use_pcp and self.impl.dcp_world_size > 1:

584 + assert lse is not None and self.dcp_manager is not None

585 + seq_lens = (

586 + attn_metadata.decode.seq_lens

587 + if attn_metadata.decode is not None

588 + else cast(torch.Tensor, attn_metadata.seq_lens)[ # type: ignore[attr-defined]

589 + : attn_metadata.num_decodes

590 + ]

591 + )

592 + query_start_loc = attn_metadata.query_start_loc[

593 + : attn_metadata.num_decodes + 1

594 + ]

595 + attn_out = self.dcp_manager.combine(

596 + attn_out,

597 + lse,

598 + seq_lens=seq_lens,

599 + query_start_loc=query_start_loc,

600 + )

601 + attn_out = finalize_mla_pcp_decode(attn_out, self.num_heads)

602 +

603 # NOTE(woosuk): While the below does not need to be in the eager region,

604 # we put it here to avoid copying the attention output. Move this back to the

605 # captured region once forward_mqa supports `out` argument.

@@ -531,3 +612,5 @@ class DeepseekV32Attention(MLAAttention):

612 .transpose(0, 1)

613 )

614 torch.bmm(x, self.W_UV, out=out)

615 + if self.use_pcp and num_actual < output.shape[0]:

616 + output[num_actual:].zero_()

vllm/models/deepseek_v32/common/kernels.py

+130 -85

@@ -103,12 +103,16 @@ def _fused_norm_rope_kernel(

103 kv_stride,

104 kv_rms_norm_w_ptr,

105 kv_rms_eps,

106 + kv_out_ptr,

107 + kv_out_stride,

108 KV_DIM: tl.constexpr,

109 # KV RoPE

110 kpe_ptr,

111 kpe_stride,

112 kpe_rope_cos_sin_cache_ptr,

113 kpe_rope_cos_sin_cache_stride,

114 + kpe_out_ptr,

115 + kpe_out_stride,

116 KPE_HALF_ROT_DIM: tl.constexpr,

117 # Index K layer norm

118 index_k_ptr,

@@ -121,6 +125,8 @@ def _fused_norm_rope_kernel(

125 # Index K RoPE

126 index_k_rope_cos_sin_cache_ptr,

127 index_k_rope_cos_sin_cache_stride,

128 + index_k_out_ptr,

129 + index_k_out_stride,

130 INDEX_K_HALF_ROT_DIM: tl.constexpr,

131 # Cache params (shared by indexer K and MLA)

132 slot_mapping_ptr,

@@ -174,10 +180,9 @@ def _fused_norm_rope_kernel(

180 return

181

182 if slot_mapping_ptr is None:

177 - # Memory profiling run.

178 - return

179 - slot_idx = tl.load(slot_mapping_ptr + tok_idx)

180 - if slot_idx < 0:

183 + if kv_out_ptr is None and kpe_out_ptr is None and index_k_out_ptr is None:

184 + return

185 + elif tl.load(slot_mapping_ptr + tok_idx) < 0:

186 # Padding

187 return

188

@@ -217,63 +222,81 @@ def _fused_norm_rope_kernel(

222 r1 = x1 * cos - x2 * sin

223 r2 = x2 * cos + x1 * sin

224

220 - # MLA concat_and_cache: write [kv_c_normed, k_pe_roped] to cache.

221 - if mla_cache_entry_stride == 0:

222 - return

223 -

224 - mla_block_size = mla_cache_block_stride // mla_cache_entry_stride

225 - mla_block_idx = slot_idx // mla_block_size

226 - mla_block_off = slot_idx % mla_block_size

227 -

228 - if MLA_CACHE_DS_MLA:

229 - # fp8_ds_mla layout (DeepSeek-V3.2, KV_DIM == 512): per-128-element

230 - # tile of the NoPE is dynamically quantized to fp8 with its own

231 - # float32 scale; the RoPE tail is stored unquantized in bf16.

232 - # bytes [0, KV_DIM) : KV_DIM fp8 NoPE values

233 - # bytes [KV_DIM, KV_DIM + 16) : MLA_NUM_TILES float32 scales

234 - # bytes [KV_DIM + 16, ...) : 2 * KPE_HALF_ROT_DIM bf16 RoPE

235 - # mla_cache_block_stride / mla_cache_entry_stride are byte strides

236 - # (mla_cache_ptr is the 1-byte fp8 view of the uint8 cache).

237 - byte_base = (

238 - mla_block_idx * mla_cache_block_stride

239 - + mla_block_off * mla_cache_entry_stride

225 + # PCP materializes K for the cross-rank gather, then inserts it into cache.

226 + if kv_out_ptr is not None:

227 + tl.store(kv_out_ptr + tok_idx * kv_out_stride + kv_block, kv_c)

228 + if kpe_out_ptr is not None:

229 + tl.store(

230 + kpe_out_ptr + tok_idx * kpe_out_stride + dim_off * 2,

231 + r1.to(kpe_out_ptr.dtype.element_ty),

232 )

241 - kv_2d = tl.reshape(kv_c, (MLA_NUM_TILES, MLA_TILE_DIM))

242 - tile_amax = tl.max(tl.abs(kv_2d), axis=1, keep_dims=True)

243 - # scale = amax / 448 (fp8 e4m3 max), matching the reference

244 - # concat_and_cache_ds_mla kernel; floored to FLT_MIN.

245 - tile_scale = tl.maximum(tile_amax * (1.0 / 448.0), 1.1754944e-38)

246 - kv_c_fp8 = tl.reshape((kv_2d / tile_scale).to(tl.float8e4nv), (KV_DIM,))

247 - tl.store(mla_cache_ptr + byte_base + kv_block, kv_c_fp8)

248 - tile_off = tl.arange(0, MLA_NUM_TILES)

233 tl.store(

250 - mla_cache_ds_scale_ptr + byte_base // 4 + KV_DIM // 4 + tile_off,

251 - tl.reshape(tile_scale, (MLA_NUM_TILES,)),

234 + kpe_out_ptr + tok_idx * kpe_out_stride + dim_off * 2 + 1,

235 + r2.to(kpe_out_ptr.dtype.element_ty),

236 )

253 - rope_dst = mla_cache_ds_rope_ptr + byte_base // 2 + (KV_DIM // 2 + 8)

254 - tl.store(rope_dst + dim_off * 2, r1.to(tl.bfloat16))

255 - tl.store(rope_dst + dim_off * 2 + 1, r2.to(tl.bfloat16))

256 - return

237

258 - dst = (

259 - mla_cache_ptr

260 - + mla_block_idx * mla_cache_block_stride

261 - + mla_block_off * mla_cache_entry_stride

262 - )

263 - # kv_c_normed (KV_DIM elements)

264 - if MLA_CACHE_FP8:

265 - scale = tl.load(mla_cache_scale_ptr)

266 - kv_c_fp8 = (kv_c.to(tl.float32) / scale).to(tl.float8e4nv)

267 - tl.store(dst + kv_block, kv_c_fp8)

268 - else:

269 - tl.store(dst + kv_block, kv_c)

270 - # k_pe_roped (from registers, interleaved layout)

271 - if MLA_CACHE_FP8:

272 - tl.store(dst + KV_DIM + dim_off * 2, (r1 / scale).to(tl.float8e4nv))

273 - tl.store(dst + KV_DIM + dim_off * 2 + 1, (r2 / scale).to(tl.float8e4nv))

274 - else:

275 - tl.store(dst + KV_DIM + dim_off * 2, r1)

276 - tl.store(dst + KV_DIM + dim_off * 2 + 1, r2)

238 + if slot_mapping_ptr is not None:

239 + # MLA concat_and_cache: write [kv_c_normed, k_pe_roped] to cache.

240 + if mla_cache_entry_stride == 0:

241 + return

242 +

243 + slot_idx = tl.load(slot_mapping_ptr + tok_idx)

244 + mla_block_size = mla_cache_block_stride // mla_cache_entry_stride

245 + mla_block_idx = slot_idx // mla_block_size

246 + mla_block_off = slot_idx % mla_block_size

247 +

248 + if MLA_CACHE_DS_MLA:

249 + # fp8_ds_mla layout (DeepSeek-V3.2, KV_DIM == 512): per-128-element

250 + # tile of the NoPE is dynamically quantized to fp8 with its own

251 + # float32 scale; the RoPE tail is stored unquantized in bf16.

252 + # bytes [0, KV_DIM) : KV_DIM fp8 NoPE values

253 + # bytes [KV_DIM, KV_DIM + 16) : MLA_NUM_TILES float32 scales

254 + # bytes [KV_DIM + 16, ...) : 2 * KPE_HALF_ROT_DIM bf16 RoPE

255 + # mla_cache_block_stride / mla_cache_entry_stride are byte strides

256 + # (mla_cache_ptr is the 1-byte fp8 view of the uint8 cache).

257 + byte_base = (

258 + mla_block_idx * mla_cache_block_stride

259 + + mla_block_off * mla_cache_entry_stride

260 + )

261 + kv_2d = tl.reshape(kv_c, (MLA_NUM_TILES, MLA_TILE_DIM))

262 + tile_amax = tl.max(tl.abs(kv_2d), axis=1, keep_dims=True)

263 + # scale = amax / 448 (fp8 e4m3 max), matching the reference

264 + # concat_and_cache_ds_mla kernel; floored to FLT_MIN.

265 + tile_scale = tl.maximum(tile_amax * (1.0 / 448.0), 1.1754944e-38)

266 + kv_c_fp8 = tl.reshape((kv_2d / tile_scale).to(tl.float8e4nv), (KV_DIM,))

267 + tl.store(mla_cache_ptr + byte_base + kv_block, kv_c_fp8)

268 + tile_off = tl.arange(0, MLA_NUM_TILES)

269 + tl.store(

270 + mla_cache_ds_scale_ptr + byte_base // 4 + KV_DIM // 4 + tile_off,

271 + tl.reshape(tile_scale, (MLA_NUM_TILES,)),

272 + )

273 + rope_dst = mla_cache_ds_rope_ptr + byte_base // 2 + (KV_DIM // 2 + 8)

274 + tl.store(rope_dst + dim_off * 2, r1.to(tl.bfloat16))

275 + tl.store(rope_dst + dim_off * 2 + 1, r2.to(tl.bfloat16))

276 + return

277 +

278 + dst = (

279 + mla_cache_ptr

280 + + mla_block_idx * mla_cache_block_stride

281 + + mla_block_off * mla_cache_entry_stride

282 + )

283 + # kv_c_normed (KV_DIM elements)

284 + if MLA_CACHE_FP8:

285 + scale = tl.load(mla_cache_scale_ptr)

286 + kv_c_fp8 = (kv_c.to(tl.float32) / scale).to(tl.float8e4nv)

287 + tl.store(dst + kv_block, kv_c_fp8)

288 + else:

289 + tl.store(dst + kv_block, kv_c)

290 + # k_pe_roped (from registers, interleaved layout)

291 + if MLA_CACHE_FP8:

292 + tl.store(dst + KV_DIM + dim_off * 2, (r1 / scale).to(tl.float8e4nv))

293 + tl.store(

294 + dst + KV_DIM + dim_off * 2 + 1,

295 + (r2 / scale).to(tl.float8e4nv),

296 + )

297 + else:

298 + tl.store(dst + KV_DIM + dim_off * 2, r1)

299 + tl.store(dst + KV_DIM + dim_off * 2 + 1, r2)

300 elif pid == 0:

301 if not HAS_INDEXER:

302 # Shared layer: no indexer K to process.

@@ -354,20 +377,27 @@ def _fused_norm_rope_kernel(

377 roped = normed * cos_full + sign * normed_partner * sin_full

378 result = tl.where(in_rope, roped, normed)

379

357 - # 3. FP8 quantize + cache write from registers.

358 - # No need to write back to index_k_ptr — the only consumer

359 - # (sparse_attn_indexer) reads from the cache, not index_k.

360 - _fp8_quant_and_cache_write(

361 - result,

362 - index_k_mask,

363 - slot_idx,

364 - indexer_cache_ptr,

365 - indexer_cache_scale_ptr,

366 - indexer_cache_block_size,

367 - indexer_cache_stride,

368 - index_k_block,

369 - INDEX_K_DIM,

370 - )

380 + if index_k_out_ptr is not None:

381 + tl.store(

382 + index_k_out_ptr + tok_idx * index_k_out_stride + index_k_block,

383 + result.to(index_k_out_ptr.dtype.element_ty),

384 + mask=index_k_mask,

385 + )

386 +

387 + # PCP inserts index K after gathering; other paths write it directly.

388 + if indexer_cache_ptr is not None and slot_mapping_ptr is not None:

389 + slot_idx = tl.load(slot_mapping_ptr + tok_idx)

390 + _fp8_quant_and_cache_write(

391 + result,

392 + index_k_mask,

393 + slot_idx,

394 + indexer_cache_ptr,

395 + indexer_cache_scale_ptr,

396 + indexer_cache_block_size,

397 + indexer_cache_stride,

398 + index_k_block,

399 + INDEX_K_DIM,

400 + )

401

402

403 def fused_norm_rope(

@@ -395,6 +425,9 @@ def fused_norm_rope(

425 has_indexer: bool = True,

426 index_rope_interleave: bool = False,

427 q_c_out: torch.Tensor | None = None,

428 + kv_c_out: torch.Tensor | None = None,

429 + k_pe_out: torch.Tensor | None = None,

430 + index_k_out: torch.Tensor | None = None,

431 ) -> torch.Tensor:

432 assert positions.ndim == 1

433 assert q_c.ndim == 2

@@ -420,6 +453,10 @@ def fused_norm_rope(

453 assert index_k_rope_cos_sin_cache is not None

454 index_k_dim = index_k.shape[-1]

455 topk = topk_indices_buffer.shape[-1]

456 + if indexer_k_cache is not None or mla_kv_cache is not None:

457 + assert slot_mapping is not None

458 + else:

459 + slot_mapping = None

460

461 # --- Indexer K cache setup ---

462 if indexer_k_cache is not None:

@@ -430,17 +467,9 @@ def fused_norm_rope(

467 if indexer_k_cache.dtype == torch.uint8:

468 indexer_k_cache = indexer_k_cache.view(torch.float8_e4m3fn)

469 else:

433 - # No indexer cache (shared layer / MLA-only fusion). Use dummies but

434 - # KEEP the caller's slot_mapping so the MLA write (pid 1) still runs.

435 - idx_cache_scale_view = torch.empty(0, dtype=torch.float32, device=device)

436 - indexer_k_cache = torch.empty(0, dtype=torch.float8_e4m3fn, device=device)

470 + idx_cache_scale_view = None

471 idx_cache_block_size = 1

438 - idx_cache_stride = 1

439 - if mla_kv_cache is None:

440 - # Pure profiling run (no caches at all): skip all per-token writes.

441 - slot_mapping = torch.full(

442 - (num_tokens,), -1, dtype=torch.int64, device=device

443 - )

472 + idx_cache_stride = 0

473

474 # --- MLA KV cache setup ---

475 mla_cache_ds_mla = mla_kv_cache_dtype == "fp8_ds_mla"

@@ -469,15 +498,26 @@ def fused_norm_rope(

498 if mla_k_scale is None:

499 mla_k_scale = torch.ones(1, dtype=torch.float32, device=device)

500 else:

472 - # Dummy values — pid 2 will skip the MLA cache write because

473 - # slot_mapping is all -1.

501 + # Dummy cache values; a zero entry stride disables the cache write.

502 mla_kv_cache = torch.empty(0, dtype=torch.bfloat16, device=device)

503 mla_block_stride = 0

504 mla_entry_stride = 0

477 - mla_k_scale = torch.ones(1, dtype=torch.float32, device=device)

505 + mla_k_scale = _dummy((1,), torch.float32, device)

506

507 if q_c_out is None:

508 q_c_out = torch.empty_like(q_c)

509 + kv_c_out_stride = 0

510 + k_pe_out_stride = 0

511 + index_k_out_stride = 0

512 + if kv_c_out is not None:

513 + assert kv_c_out.shape == kv_c.shape

514 + kv_c_out_stride = kv_c_out.stride(0)

515 + if k_pe_out is not None:

516 + assert k_pe_out.shape == k_pe.shape

517 + k_pe_out_stride = k_pe_out.stride(0)

518 + if index_k_out is not None:

519 + assert index_k_out.shape == index_k.shape

520 + index_k_out_stride = index_k_out.stride(0)

521 use_pdl = current_platform.is_arch_support_pdl()

522 _fused_norm_rope_kernel[(num_tokens, 4)](

523 positions,

@@ -495,12 +535,16 @@ def fused_norm_rope(

535 kv_c.stride(0),

536 kv_rms_norm_w,

537 kv_rms_eps,

538 + kv_c_out,

539 + kv_c_out_stride,

540 kv_dim,

541 # KV RoPE

542 k_pe,

543 k_pe.stride(0),

544 k_rope_cos_sin_cache,

545 k_rope_cos_sin_cache.stride(0),

546 + k_pe_out,

547 + k_pe_out_stride,

548 k_rope_cos_sin_cache.shape[-1] // 2,

549 # Index K layer norm + RoPE + FP8 quant

550 index_k,

@@ -512,6 +556,8 @@ def fused_norm_rope(

556 triton.next_power_of_2(index_k_dim),

557 index_k_rope_cos_sin_cache,

558 index_k_rope_cos_sin_cache.stride(0),

559 + index_k_out,

560 + index_k_out_stride,

561 index_k_rope_cos_sin_cache.shape[-1] // 2,

562 # Cache params

563 slot_mapping,

@@ -793,7 +839,6 @@ def fused_q(

839 assert ql_nope.ndim == 3

840 assert ql_nope.shape[:2] == q_pe.shape[:2]

841 assert q_scale.dtype == torch.float32 and q_scale.numel() == 1

796 -

842 num_tokens = positions.shape[0]

843 num_q_heads = q_pe.shape[1]

844 # Grid's 3rd dim must cover the MQA-pack heads (pid 0/2 iterate 2 heads

vllm/v1/attention/ops/rocm_aiter_mla_sparse.py

+2 -2

@@ -682,7 +682,7 @@ def rocm_aiter_sparse_attn_indexer_fake(

682 k_cache_prefix: LayerNameType,

683 kv_cache: torch.Tensor,

684 q_fp8: torch.Tensor,

685 - k: torch.Tensor,

685 + k: torch.Tensor | None,

686 weights: torch.Tensor,

687 quant_block_size: int,

688 scale_fmt: str | None,

@@ -702,7 +702,7 @@ def rocm_aiter_sparse_attn_indexer(

702 k_cache_prefix: LayerNameType,

703 kv_cache: torch.Tensor,

704 q_fp8: torch.Tensor,

705 - k: torch.Tensor,

705 + k: torch.Tensor | None,

706 weights: torch.Tensor,

707 quant_block_size: int,

708 scale_fmt: str | None,