@tduncan / Nvidia-Optimization-Stack-VLLM / commits / eb33ff34dd

[ROCm][Perf] DSv4 two-stage compressor kernel for HCA prefill (#47718)

Signed-off-by: kliuae <kuanfu.liu@embeddedllm.com>

kliuae committed Jul 16, 2026 at 07:31 UTC eb33ff34dd6501f23f5be92a5c842d34ab05d7fd
3 files changed +528
tests/kernels/test_compressor_kv_cache.py
+130
@@ -9,6 +9,7 @@ These tests cover:
9 C) Indexer: head_dim=128 (all FP8), quant_block=128
10 D) DeepseekV4 Attention magnitude range: correctness across small/large values
11 E) Indexer fused Triton kernel: compress+norm+rope+quant+insert
12 + F) Indexer fused two-stage Triton kernel: head=512 cr>=128 (no-overlap)
13 """
14
15 import math
@@ -24,7 +25,9 @@ from vllm.models.deepseek_v4.common.ops import (
25 from vllm.models.deepseek_v4.common.ops.fused_compress_quant_cache import (
26 _fused_kv_compress_norm_rope_insert_indexer_attn,
27 _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn,
28 + _launch_two_stage_sparse_attn_compressor,
29 )
30 +from vllm.platforms import current_platform
31
32 from .test_fused_indexer_q_rope_quant import quantize_to_mxfp4
33
@@ -816,3 +819,130 @@ def test_cutedsl_full_cache_store(compress_ratio: int, store_fp8: bool):
819 torch.testing.assert_close(actual.float(), ref_fp8.float(), rtol=0.0, atol=0.3)
820 else:
821 torch.testing.assert_close(actual.float(), ref.float(), rtol=3e-2, atol=3e-2)
822 +
823 +
824 +# ── Test F: DeepseekV4 Attention two-stage split compressor (Triton) ─────────
825 +#
826 +# Same full pipeline as Test E (state-cache gather -> softmax-weighted compress
827 +# -> RMSNorm -> GPT-J RoPE -> quant -> paged insert), but for the head=512
828 +# fp8_ds_mla layout via the two-stage split
829 +
830 +
831 +@pytest.mark.skipif(
832 + not current_platform.is_rocm(),
833 + reason="two-stage split compressor is only enabled for ROCm at the moment",
834 +)
835 +@pytest.mark.parametrize("num_tokens", [1, 4, 8, 17])
836 +@pytest.mark.parametrize("kv_block_size", [16, 64])
837 +def test_fused_kv_insert_split(num_tokens: int, kv_block_size: int):
838 + """Two-stage split compress+norm+rope+quant+insert for the head=512 KV cache."""
839 + HEAD_DIM = 512
840 + NOPE_DIM = 448
841 + ROPE_DIM = 64
842 + HEAD_BYTES = 584 # 448 fp8 + 128 bf16 + 8 uint8 scale
843 + RMS_EPS = 1e-6
844 + FP8_MAX = 448.0
845 + QUANT_BLOCK = 64
846 + TOKEN_STRIDE = 576
847 + SCALE_DIM = 8
848 + STATE_BLOCK_SIZE = 8 # CompressorStateCache block_size for cr=128
849 +
850 + device = "cuda"
851 + torch.manual_seed(42)
852 + compress_ratio = 128
853 + overlap = 0 # no overlap for cr=128
854 + coff = 1 + overlap
855 +
856 + num_pages = (compress_ratio * num_tokens - 1) // STATE_BLOCK_SIZE + 2
857 + state_cache = torch.randn(
858 + num_pages,
859 + STATE_BLOCK_SIZE,
860 + 2 * coff * HEAD_DIM, # kv_state + score_state
861 + dtype=torch.float32,
862 + device=device,
863 + )
864 + block_table = torch.arange(num_pages, dtype=torch.int32, device=device).unsqueeze(0)
865 + token_to_req = torch.zeros(num_tokens, dtype=torch.int32, device=device)
866 + slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device)
867 + positions = torch.arange(
868 + compress_ratio - 1,
869 + compress_ratio * num_tokens,
870 + compress_ratio,
871 + dtype=torch.int64,
872 + device=device,
873 + )
874 + rms_weight = torch.randn(HEAD_DIM, dtype=torch.bfloat16, device=device)
875 + cos_sin_cache = torch.randn(
876 + compress_ratio * num_tokens, ROPE_DIM, dtype=torch.float32, device=device
877 + )
878 +
879 + kv_n_blocks = (num_tokens + kv_block_size - 1) // kv_block_size + 1
880 + kv_cache = torch.zeros(
881 + kv_n_blocks, kv_block_size, HEAD_BYTES, dtype=torch.uint8, device=device
882 + )
883 + compress_scratch = torch.empty(
884 + num_tokens, HEAD_DIM, dtype=torch.float32, device=device
885 + )
886 +
887 + _launch_two_stage_sparse_attn_compressor(
888 + state_cache,
889 + token_to_req,
890 + positions,
891 + slot_mapping,
892 + block_table,
893 + STATE_BLOCK_SIZE,
894 + coff * HEAD_DIM,
895 + compress_ratio,
896 + cos_sin_cache,
897 + kv_cache,
898 + slot_mapping,
899 + rms_weight,
900 + RMS_EPS,
901 + QUANT_BLOCK,
902 + TOKEN_STRIDE,
903 + SCALE_DIM,
904 + HEAD_DIM,
905 + ROPE_DIM,
906 + num_tokens,
907 + compress_scratch,
908 + )
909 +
910 + # PyTorch reference: compress -> RMSNorm -> GPT-J RoPE (pre-quant bf16 row).
911 + ref = _reference_kv_compress_norm_rope(
912 + state_cache,
913 + block_table,
914 + positions,
915 + rms_weight,
916 + cos_sin_cache,
917 + compress_ratio,
918 + overlap,
919 + rms_eps=RMS_EPS,
920 + fp8_max=FP8_MAX,
921 + return_full_cache=True,
922 + ) # [num_tokens, HEAD_DIM] bf16
923 +
924 + # Dequant + gather the fp8_ds_mla cache back to bf16 (Test B op).
925 + out = torch.zeros(1, num_tokens, HEAD_DIM, dtype=torch.bfloat16, device=device)
926 + seq_lens = torch.tensor([num_tokens], dtype=torch.int32, device=device)
927 + gather_block_table = torch.arange(
928 + kv_n_blocks, dtype=torch.int32, device=device
929 + ).unsqueeze(0)
930 + dequantize_and_gather_k_cache(
931 + out, kv_cache, seq_lens, None, gather_block_table, kv_block_size, offset=0
932 + )
933 + recovered = out[0, :num_tokens]
934 +
935 + # NoPE (first 448): FP8 quantized, expect UE8M0 error (same bound as Test A).
936 + nope_diff = (recovered[:, :NOPE_DIM].float() - ref[:, :NOPE_DIM].float()).abs()
937 + for t in range(num_tokens):
938 + _, scales = _ue8m0_reference(ref[t, :NOPE_DIM].float(), QUANT_BLOCK, FP8_MAX)
939 + max_allowed = 16.0 * scales.max().item()
940 + token_diff = nope_diff[t].max().item()
941 + assert token_diff <= max_allowed, (
942 + f"Token {t} nope diff {token_diff} exceeds max_allowed "
943 + f"{max_allowed} (scale={scales.max().item()})"
944 + )
945 +
946 + # RoPE (last 64): stored as bf16. The kernel recomputes the rotation, so it
947 + # is bf16-close to the reference rather than bit-exact (cf. test_cutedsl).
948 + torch.testing.assert_close(recovered[:, NOPE_DIM:], ref[:, NOPE_DIM:])
vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py
+355
@@ -19,6 +19,7 @@ even/odd halves, producing (N_QUANT_BLOCKS, MXFP4_BLOCK/2) packed nibbles
19 and N_QUANT_BLOCKS ue8m0 bytes.
20 """
21
22 +from functools import lru_cache
23 from typing import Any
24
25 import torch
@@ -296,6 +297,360 @@ def _fused_kv_compress_norm_rope_insert_sparse_attn(
297 tl.store(bf16_ptr + rope_local, result.to(tl.bfloat16), mask=is_rope)
298
299
300 +# =============================================================================
301 +# Split kernels variant of the head=512 compressor (deep cr=128 gather).
302 +# - compress gather: instead of launching one program per token, split along
303 +# the head dimension to maximize CU occupancy. The head dimension split
304 +# does not require cross-group reduction
305 +# - finalize norm rope quant store: same as the single pass kernel due to its
306 +# per-token nature
307 +# Mirrors the CUDA cutedsl split kernel where num_splits is occupancy-targeted.
308 +# Currently only tested and validated on ROCm gfx950
309 +# =============================================================================
310 +@lru_cache(maxsize=1)
311 +def _n_cu() -> int:
312 + return torch.cuda.get_device_properties(0).multi_processor_count
313 +
314 +
315 +def _pick_compress_num_splits(
316 + num_actual: int, compress_ratio: int, head_dim: int
317 +) -> int:
318 + """Occupancy-targeted column splits for the cr>=128 head=512 compressor.
319 +
320 + Sizes the per-token fan-out so (estimated computing tokens) * num_splits ~
321 + #CU, capped by head tiling at a 32-wide min tile, as a power-of-2 divisor of
322 + head_dim.
323 + """
324 + max_splits = head_dim // 32
325 + est_compute = max(1, num_actual // compress_ratio)
326 + target = -(-_n_cu() // est_compute) # ceil(#CU / est_compute)
327 + ns = 1
328 + while ns * 2 <= min(target, max_splits) and head_dim % (ns * 2) == 0:
329 + ns *= 2
330 + return ns
331 +
332 +
333 +@triton.jit
334 +def _compress_gather_split_sparse_attn(
335 + state_cache_ptr,
336 + state_cache_stride0,
337 + state_cache_stride1,
338 + positions_ptr,
339 + slot_mapping_ptr,
340 + token_to_req_indices_ptr,
341 + block_table_ptr,
342 + block_table_stride,
343 + block_size,
344 + scratch_ptr,
345 + scratch_stride,
346 + HEAD_SIZE: tl.constexpr,
347 + STATE_WIDTH: tl.constexpr,
348 + COMPRESS_RATIO: tl.constexpr,
349 + NUM_SPLITS: tl.constexpr,
350 + HEAD_TILE: tl.constexpr, # HEAD_SIZE // NUM_SPLITS
351 +):
352 + """Stage 1: per-(token, head-split) compress gather, write to fp32 scratch
353 +
354 + No-overlap gather (cr>=128) on rows [0, COMPRESS_RATIO)
355 + """
356 + pid = tl.program_id(0)
357 + token_idx = pid // NUM_SPLITS
358 + split_idx = pid % NUM_SPLITS
359 +
360 + slot_id = tl.load(slot_mapping_ptr + token_idx)
361 + if slot_id < 0:
362 + return
363 + position = tl.load(positions_ptr + token_idx)
364 + if (position + 1) % COMPRESS_RATIO != 0:
365 + return
366 + req_idx = tl.load(token_to_req_indices_ptr + token_idx)
367 +
368 + start = position - COMPRESS_RATIO + 1
369 + rows = tl.arange(0, COMPRESS_RATIO)
370 + pos = start + rows
371 + mask_pos = pos >= 0
372 + block_numbers = tl.load(
373 + block_table_ptr + req_idx * block_table_stride + pos // block_size,
374 + mask=mask_pos,
375 + other=0,
376 + ).to(tl.int64)
377 + block_offsets = pos % block_size
378 +
379 + col = split_idx * HEAD_TILE + tl.arange(0, HEAD_TILE)
380 + row_base = (
381 + state_cache_ptr
382 + + block_numbers * state_cache_stride0
383 + + block_offsets * state_cache_stride1
384 + )
385 + cmask = mask_pos[:, None]
386 +
387 + score = tl.load(
388 + row_base[:, None] + STATE_WIDTH + col[None, :],
389 + mask=cmask,
390 + other=float("-inf"),
391 + )
392 + score = tl.softmax(score, dim=0)
393 + kv = tl.load(row_base[:, None] + col[None, :], mask=cmask, other=0.0)
394 + compressed = tl.sum(kv * score, axis=0) # [HEAD_TILE] fp32
395 + tl.store(scratch_ptr + token_idx * scratch_stride + col, compressed)
396 +
397 +
398 +@triton.jit
399 +def _finalize_norm_rope_quant_store_sparse_attn(
400 + scratch_ptr,
401 + scratch_stride,
402 + positions_ptr,
403 + slot_mapping_ptr,
404 + rms_norm_weight_ptr,
405 + rms_norm_eps,
406 + cos_sin_cache_ptr,
407 + cos_sin_stride,
408 + k_cache_ptr,
409 + kv_slot_mapping_ptr,
410 + kv_cache_block_size,
411 + HEAD_SIZE: tl.constexpr,
412 + TRITON_BLOCK_SIZE: tl.constexpr,
413 + COMPRESS_RATIO: tl.constexpr,
414 + ROPE_HEAD_DIM: tl.constexpr,
415 + FP8_MAX: tl.constexpr,
416 + QUANT_BLOCK: tl.constexpr,
417 + TOKEN_STRIDE: tl.constexpr,
418 + SCALE_DIM: tl.constexpr,
419 + KV_BLOCK_STRIDE: tl.constexpr,
420 +):
421 + """Stage 2: read compressed_kv[512] from scratch buffer, then
422 + RMSNorm + FP8 quant (nope) + RoPE + bf16 store
423 + """
424 + token_idx = tl.program_id(0)
425 + slot_id = tl.load(slot_mapping_ptr + token_idx)
426 + if slot_id < 0:
427 + return
428 + position = tl.load(positions_ptr + token_idx)
429 + if (position + 1) % COMPRESS_RATIO != 0:
430 + return
431 +
432 + block = tl.arange(0, TRITON_BLOCK_SIZE)
433 + mask = block < HEAD_SIZE
434 + compressed_kv = tl.load(
435 + scratch_ptr + token_idx * scratch_stride + block, mask=mask, other=0.0
436 + )
437 +
438 + rms_w = tl.load(rms_norm_weight_ptr + block, mask=mask, other=0.0)
439 + variance = tl.sum(compressed_kv * compressed_kv, axis=0) / HEAD_SIZE
440 + rrms = tl.rsqrt(variance + rms_norm_eps)
441 + normed = compressed_kv * rrms * rms_w
442 +
443 + kv_slot_idx = tl.load(kv_slot_mapping_ptr + token_idx)
444 + if kv_slot_idx < 0:
445 + return
446 + kv_block_idx = kv_slot_idx // kv_cache_block_size
447 + kv_pos_in_block = kv_slot_idx % kv_cache_block_size
448 + cache_block_ptr = k_cache_ptr + kv_block_idx.to(tl.int64) * KV_BLOCK_STRIDE
449 + fp8_ptr = cache_block_ptr + kv_pos_in_block * TOKEN_STRIDE
450 + scale_ptr = (
451 + cache_block_ptr
452 + + kv_cache_block_size * TOKEN_STRIDE
453 + + kv_pos_in_block * SCALE_DIM
454 + )
455 +
456 + NOPE_HEAD_DIM: tl.constexpr = HEAD_SIZE - ROPE_HEAD_DIM
457 + HALF_ROPE: tl.constexpr = ROPE_HEAD_DIM // 2
458 + N_QUANT_BLOCKS: tl.constexpr = TRITON_BLOCK_SIZE // QUANT_BLOCK
459 + N_NOPE_BLOCKS: tl.constexpr = NOPE_HEAD_DIM // QUANT_BLOCK
460 + INV_FP8_MAX: tl.constexpr = 1.0 / FP8_MAX
461 +
462 + quant_input = normed.to(tl.bfloat16).to(tl.float32)
463 + quant_2d = tl.reshape(quant_input, (N_QUANT_BLOCKS, QUANT_BLOCK))
464 + block_absmax = tl.maximum(tl.max(tl.abs(quant_2d), axis=1), 1e-4)
465 + raw_scales = block_absmax * INV_FP8_MAX
466 + exponents = tl.ceil(tl.log2(raw_scales))
467 + inv_scales = tl.exp2(-exponents)
468 + x_scaled = quant_2d * tl.reshape(inv_scales, (N_QUANT_BLOCKS, 1))
469 + x_clamped = tl.clamp(x_scaled, -FP8_MAX, FP8_MAX)
470 + x_uint8 = tl.reshape(
471 + x_clamped.to(tl.float8e4nv).to(tl.uint8, bitcast=True),
472 + (TRITON_BLOCK_SIZE,),
473 + )
474 + tl.store(fp8_ptr + block, x_uint8, mask=block < NOPE_HEAD_DIM)
475 +
476 + scale_idx = tl.arange(0, N_QUANT_BLOCKS)
477 + encoded = tl.maximum(tl.minimum(exponents + 127.0, 255.0), 0.0)
478 + tl.store(
479 + scale_ptr + scale_idx, encoded.to(tl.uint8), mask=scale_idx < N_NOPE_BLOCKS
480 + )
481 + tl.store(scale_ptr + N_NOPE_BLOCKS, tl.zeros((), dtype=tl.uint8))
482 +
483 + NUM_PAIRS: tl.constexpr = TRITON_BLOCK_SIZE // 2
484 + NOPE_PAIRS: tl.constexpr = NOPE_HEAD_DIM // 2
485 + even, odd = tl.split(tl.reshape(normed, (NUM_PAIRS, 2)))
486 + pair_idx = tl.arange(0, NUM_PAIRS)
487 + rope_pair_local = pair_idx - NOPE_PAIRS
488 + is_rope_pair = rope_pair_local >= 0
489 + cs_idx = tl.maximum(rope_pair_local, 0)
490 + compressed_pos = (position // COMPRESS_RATIO) * COMPRESS_RATIO
491 + cache_base = cos_sin_cache_ptr + compressed_pos * cos_sin_stride
492 + cos_v = tl.load(cache_base + cs_idx, mask=is_rope_pair, other=1.0)
493 + sin_v = tl.load(cache_base + HALF_ROPE + cs_idx, mask=is_rope_pair, other=0.0)
494 + new_even = even * cos_v - odd * sin_v
495 + new_odd = odd * cos_v + even * sin_v
496 + result = tl.interleave(new_even, new_odd)
497 + bf16_ptr = (fp8_ptr + NOPE_HEAD_DIM).to(tl.pointer_type(tl.bfloat16))
498 + rope_local = block - NOPE_HEAD_DIM
499 + is_rope = (block >= NOPE_HEAD_DIM) & mask
500 + tl.store(bf16_ptr + rope_local, result.to(tl.bfloat16), mask=is_rope)
501 +
502 +
503 +def _launch_two_stage_sparse_attn_compressor(
504 + state_cache: torch.Tensor,
505 + token_to_req_indices: torch.Tensor,
506 + positions: torch.Tensor,
507 + slot_mapping: torch.Tensor,
508 + block_table: torch.Tensor,
509 + block_size: int,
510 + state_width: int,
511 + compress_ratio: int,
512 + cos_sin_cache: torch.Tensor,
513 + kv_cache: torch.Tensor,
514 + kv_slot_mapping: torch.Tensor,
515 + rms_norm_weight: torch.Tensor,
516 + rms_norm_eps: float,
517 + quant_block: int,
518 + token_stride: int,
519 + scale_dim: int,
520 + head_dim: int,
521 + rope_head_dim: int,
522 + num_actual: int,
523 + compress_scratch: torch.Tensor,
524 +) -> None:
525 + num_splits = _pick_compress_num_splits(num_actual, compress_ratio, head_dim)
526 + head_tile = head_dim // num_splits
527 + scratch = compress_scratch[:num_actual]
528 + _compress_gather_split_sparse_attn[(num_actual * num_splits,)](
529 + state_cache,
530 + state_cache.stride(0),
531 + state_cache.stride(1),
532 + positions,
533 + slot_mapping,
534 + token_to_req_indices,
535 + block_table,
536 + block_table.stride(0),
537 + block_size,
538 + scratch,
539 + scratch.stride(0),
540 + HEAD_SIZE=head_dim,
541 + STATE_WIDTH=state_width,
542 + COMPRESS_RATIO=compress_ratio,
543 + NUM_SPLITS=num_splits,
544 + HEAD_TILE=head_tile,
545 + )
546 + _finalize_norm_rope_quant_store_sparse_attn[(num_actual,)](
547 + scratch,
548 + scratch.stride(0),
549 + positions,
550 + slot_mapping,
551 + rms_norm_weight,
552 + rms_norm_eps,
553 + cos_sin_cache,
554 + cos_sin_cache.stride(0),
555 + kv_cache,
556 + kv_slot_mapping,
557 + kv_cache.shape[1],
558 + HEAD_SIZE=head_dim,
559 + TRITON_BLOCK_SIZE=triton.next_power_of_2(head_dim),
560 + COMPRESS_RATIO=compress_ratio,
561 + ROPE_HEAD_DIM=rope_head_dim,
562 + FP8_MAX=448.0,
563 + QUANT_BLOCK=quant_block,
564 + TOKEN_STRIDE=token_stride,
565 + SCALE_DIM=scale_dim,
566 + KV_BLOCK_STRIDE=kv_cache.stride(0),
567 + )
568 +
569 +
570 +def compress_norm_rope_store_two_stage_triton(
571 + state_cache: torch.Tensor,
572 + num_actual: int,
573 + token_to_req_indices: torch.Tensor,
574 + positions: torch.Tensor,
575 + slot_mapping: torch.Tensor,
576 + block_table: torch.Tensor,
577 + block_size: int,
578 + state_width: int,
579 + cos_sin_cache: torch.Tensor,
580 + kv_cache: torch.Tensor,
581 + k_cache_metadata: Any,
582 + pdl_kwargs: dict,
583 + head_dim: int,
584 + rope_head_dim: int,
585 + compress_ratio: int,
586 + overlap: bool,
587 + use_fp4_cache: bool,
588 + rms_norm_weight: torch.Tensor,
589 + rms_norm_eps: float,
590 + quant_block: int,
591 + token_stride: int,
592 + scale_dim: int,
593 + num_decode_tokens: int,
594 + compress_scratch: torch.Tensor,
595 +) -> None:
596 + """Two-stage split compressor dispatch for head=512 cr>=128 (no-overlap)
597 +
598 + Run the occupancy-fanned two-stage split for prefill [num_decodee_tokens:]
599 + to fill the CUs, and use the original single-pass launcher
600 + for decode [0, num_decode_tokens)
601 + """
602 + num_decodes = min(max(num_decode_tokens, 0), num_actual)
603 + num_prefills = num_actual - num_decodes
604 + if num_prefills > 0:
605 + _launch_two_stage_sparse_attn_compressor(
606 + state_cache=state_cache,
607 + token_to_req_indices=token_to_req_indices[num_decodes:],
608 + positions=positions[num_decodes:],
609 + slot_mapping=slot_mapping[num_decodes:],
610 + block_table=block_table,
611 + block_size=block_size,
612 + state_width=state_width,
613 + compress_ratio=compress_ratio,
614 + cos_sin_cache=cos_sin_cache,
615 + kv_cache=kv_cache,
616 + kv_slot_mapping=k_cache_metadata.slot_mapping[num_decodes:],
617 + rms_norm_weight=rms_norm_weight,
618 + rms_norm_eps=rms_norm_eps,
619 + quant_block=quant_block,
620 + token_stride=token_stride,
621 + scale_dim=scale_dim,
622 + head_dim=head_dim,
623 + rope_head_dim=rope_head_dim,
624 + num_actual=num_prefills,
625 + compress_scratch=compress_scratch,
626 + )
627 + if num_decodes > 0:
628 + compress_norm_rope_store_triton(
629 + state_cache=state_cache,
630 + num_actual=num_decodes,
631 + token_to_req_indices=token_to_req_indices,
632 + positions=positions,
633 + slot_mapping=slot_mapping,
634 + block_table=block_table,
635 + block_size=block_size,
636 + state_width=state_width,
637 + cos_sin_cache=cos_sin_cache,
638 + kv_cache=kv_cache,
639 + k_cache_metadata=k_cache_metadata,
640 + pdl_kwargs=pdl_kwargs,
641 + head_dim=head_dim,
642 + rope_head_dim=rope_head_dim,
643 + compress_ratio=compress_ratio,
644 + overlap=overlap,
645 + use_fp4_cache=use_fp4_cache,
646 + rms_norm_weight=rms_norm_weight,
647 + rms_norm_eps=rms_norm_eps,
648 + quant_block=quant_block,
649 + token_stride=token_stride,
650 + scale_dim=scale_dim,
651 + )
652 +
653 +
654 # =============================================================================
655 # Indexer path (head=128, all FP8, single quant block)
656 # =============================================================================
vllm/models/deepseek_v4/compressor.py
+43
@@ -14,6 +14,7 @@ from vllm.model_executor.layers.layernorm import RMSNorm
14 from vllm.model_executor.layers.linear import MergedColumnParallelLinear
15 from vllm.models.deepseek_v4.common.ops.fused_compress_quant_cache import (
16 compress_norm_rope_store_triton,
17 + compress_norm_rope_store_two_stage_triton,
18 )
19 from vllm.models.deepseek_v4.common.ops.fused_indexer_q import MXFP4_BLOCK_SIZE
20 from vllm.models.deepseek_v4.common.ops.save_partial_states import (
@@ -27,6 +28,7 @@ from vllm.v1.attention.backend import (
28 CommonAttentionMetadata,
29 MultipleOf,
30 )
31 +from vllm.v1.attention.backends.utils import split_decodes_and_prefills
32 from vllm.v1.kv_cache_interface import (
33 KVCacheSpec,
34 MLAAttentionSpec,
@@ -34,6 +36,12 @@ from vllm.v1.kv_cache_interface import (
36 )
37
38
39 +def _prefer_two_stage_compressor() -> bool:
40 + # Platforms that favor the triton variant of two-stage compressor split.
41 + # Currently only tested on ROCm
42 + return current_platform.is_rocm()
43 +
44 +
45 class CompressorBackend(AttentionBackend):
46 def __init__(self):
47 super().__init__()
@@ -81,6 +89,7 @@ class CompressorMetadata:
89 block_size: int
90
91 token_to_req_indices: torch.Tensor | None = None # [num_tokens]
92 + num_decode_tokens: int | None = None
93
94
95 class CompressorMetadataBuilder(AttentionMetadataBuilder):
@@ -107,11 +116,17 @@ class CompressorMetadataBuilder(AttentionMetadataBuilder):
116 token_to_req_indices = common_attn_metadata.token_to_req_indices(
117 self.token_to_req_indices
118 )
119 + num_decode_tokens = None
120 + if _prefer_two_stage_compressor():
121 + _, _, num_decode_tokens, _ = split_decodes_and_prefills(
122 + common_attn_metadata, decode_threshold=1
123 + )
124 return CompressorMetadata(
125 block_table=common_attn_metadata.block_table_tensor.clamp_(min=0),
126 slot_mapping=common_attn_metadata.slot_mapping,
127 block_size=self.block_size,
128 token_to_req_indices=token_to_req_indices,
129 + num_decode_tokens=num_decode_tokens,
130 )
131
132
@@ -213,6 +228,25 @@ class DeepseekCompressor(nn.Module):
228 self.overlap = compress_ratio == 4
229 self.coff = 1 + self.overlap
230
231 + # The head=512 cr>=128 no-overlap deep gather uses the two-stage
232 + # compressor, which needs an fp32 scratch [max_batched, 512] for
233 + # the intermediate compressed_kv.
234 + # Currently only tested on ROCm
235 + self._use_two_stage_fused_compressor = (
236 + _prefer_two_stage_compressor() and head_dim == 512 and not self.overlap
237 + )
238 + self.max_num_batched_tokens = (
239 + vllm_config.scheduler_config.max_num_batched_tokens
240 + )
241 + self._compress_scratch: torch.Tensor | None = None
242 + if self._use_two_stage_fused_compressor:
243 + self._compress_scratch = torch.empty(
244 + self.max_num_batched_tokens,
245 + self.head_dim,
246 + dtype=torch.float32,
247 + device=self.device,
248 + )
249 +
250 state_dtype = torch.float32
251 self.ape = nn.Parameter(
252 torch.empty(
@@ -364,6 +398,15 @@ class DeepseekCompressor(nn.Module):
398 store_full_fp8=store_full_fp8,
399 fp8_scale=fp8_scale,
400 )
401 + elif self._use_two_stage_fused_compressor:
402 + # head=512 cr>=128 (no overlap): two-pass split compressor on the
403 + # prefill suffix, single-pass on the decode prefix.
404 + assert state_metadata.num_decode_tokens is not None
405 + compress_norm_rope_store_fn = compress_norm_rope_store_two_stage_triton
406 + extra_kwargs = {
407 + "num_decode_tokens": state_metadata.num_decode_tokens,
408 + "compress_scratch": self._compress_scratch,
409 + }
410 else:
411 # Indexer path (head_dim == 128) or non-CUDA GPUs (AMD, XPU, etc.).
412 compress_norm_rope_store_fn = compress_norm_rope_store_triton