diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index 74dc01472a..b8f1cb8bfa 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -9,6 +9,7 @@ These tests cover: C) Indexer: head_dim=128 (all FP8), quant_block=128 D) DeepseekV4 Attention magnitude range: correctness across small/large values E) Indexer fused Triton kernel: compress+norm+rope+quant+insert + F) Indexer fused two-stage Triton kernel: head=512 cr>=128 (no-overlap) """ import math @@ -24,7 +25,9 @@ from vllm.models.deepseek_v4.common.ops import ( from vllm.models.deepseek_v4.common.ops.fused_compress_quant_cache import ( _fused_kv_compress_norm_rope_insert_indexer_attn, _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn, + _launch_two_stage_sparse_attn_compressor, ) +from vllm.platforms import current_platform from .test_fused_indexer_q_rope_quant import quantize_to_mxfp4 @@ -816,3 +819,130 @@ def test_cutedsl_full_cache_store(compress_ratio: int, store_fp8: bool): torch.testing.assert_close(actual.float(), ref_fp8.float(), rtol=0.0, atol=0.3) else: torch.testing.assert_close(actual.float(), ref.float(), rtol=3e-2, atol=3e-2) + + +# ── Test F: DeepseekV4 Attention two-stage split compressor (Triton) ───────── +# +# Same full pipeline as Test E (state-cache gather -> softmax-weighted compress +# -> RMSNorm -> GPT-J RoPE -> quant -> paged insert), but for the head=512 +# fp8_ds_mla layout via the two-stage split + + +@pytest.mark.skipif( + not current_platform.is_rocm(), + reason="two-stage split compressor is only enabled for ROCm at the moment", +) +@pytest.mark.parametrize("num_tokens", [1, 4, 8, 17]) +@pytest.mark.parametrize("kv_block_size", [16, 64]) +def test_fused_kv_insert_split(num_tokens: int, kv_block_size: int): + """Two-stage split compress+norm+rope+quant+insert for the head=512 KV cache.""" + HEAD_DIM = 512 + NOPE_DIM = 448 + ROPE_DIM = 64 + HEAD_BYTES = 584 # 448 fp8 + 128 bf16 + 8 uint8 scale + RMS_EPS = 1e-6 + FP8_MAX = 448.0 + QUANT_BLOCK = 64 + TOKEN_STRIDE = 576 + SCALE_DIM = 8 + STATE_BLOCK_SIZE = 8 # CompressorStateCache block_size for cr=128 + + device = "cuda" + torch.manual_seed(42) + compress_ratio = 128 + overlap = 0 # no overlap for cr=128 + coff = 1 + overlap + + num_pages = (compress_ratio * num_tokens - 1) // STATE_BLOCK_SIZE + 2 + state_cache = torch.randn( + num_pages, + STATE_BLOCK_SIZE, + 2 * coff * HEAD_DIM, # kv_state + score_state + dtype=torch.float32, + device=device, + ) + block_table = torch.arange(num_pages, dtype=torch.int32, device=device).unsqueeze(0) + token_to_req = torch.zeros(num_tokens, dtype=torch.int32, device=device) + slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device) + positions = torch.arange( + compress_ratio - 1, + compress_ratio * num_tokens, + compress_ratio, + dtype=torch.int64, + device=device, + ) + rms_weight = torch.randn(HEAD_DIM, dtype=torch.bfloat16, device=device) + cos_sin_cache = torch.randn( + compress_ratio * num_tokens, ROPE_DIM, dtype=torch.float32, device=device + ) + + kv_n_blocks = (num_tokens + kv_block_size - 1) // kv_block_size + 1 + kv_cache = torch.zeros( + kv_n_blocks, kv_block_size, HEAD_BYTES, dtype=torch.uint8, device=device + ) + compress_scratch = torch.empty( + num_tokens, HEAD_DIM, dtype=torch.float32, device=device + ) + + _launch_two_stage_sparse_attn_compressor( + state_cache, + token_to_req, + positions, + slot_mapping, + block_table, + STATE_BLOCK_SIZE, + coff * HEAD_DIM, + compress_ratio, + cos_sin_cache, + kv_cache, + slot_mapping, + rms_weight, + RMS_EPS, + QUANT_BLOCK, + TOKEN_STRIDE, + SCALE_DIM, + HEAD_DIM, + ROPE_DIM, + num_tokens, + compress_scratch, + ) + + # PyTorch reference: compress -> RMSNorm -> GPT-J RoPE (pre-quant bf16 row). + ref = _reference_kv_compress_norm_rope( + state_cache, + block_table, + positions, + rms_weight, + cos_sin_cache, + compress_ratio, + overlap, + rms_eps=RMS_EPS, + fp8_max=FP8_MAX, + return_full_cache=True, + ) # [num_tokens, HEAD_DIM] bf16 + + # Dequant + gather the fp8_ds_mla cache back to bf16 (Test B op). + out = torch.zeros(1, num_tokens, HEAD_DIM, dtype=torch.bfloat16, device=device) + seq_lens = torch.tensor([num_tokens], dtype=torch.int32, device=device) + gather_block_table = torch.arange( + kv_n_blocks, dtype=torch.int32, device=device + ).unsqueeze(0) + dequantize_and_gather_k_cache( + out, kv_cache, seq_lens, None, gather_block_table, kv_block_size, offset=0 + ) + recovered = out[0, :num_tokens] + + # NoPE (first 448): FP8 quantized, expect UE8M0 error (same bound as Test A). + nope_diff = (recovered[:, :NOPE_DIM].float() - ref[:, :NOPE_DIM].float()).abs() + for t in range(num_tokens): + _, scales = _ue8m0_reference(ref[t, :NOPE_DIM].float(), QUANT_BLOCK, FP8_MAX) + max_allowed = 16.0 * scales.max().item() + token_diff = nope_diff[t].max().item() + assert token_diff <= max_allowed, ( + f"Token {t} nope diff {token_diff} exceeds max_allowed " + f"{max_allowed} (scale={scales.max().item()})" + ) + + # RoPE (last 64): stored as bf16. The kernel recomputes the rotation, so it + # is bf16-close to the reference rather than bit-exact (cf. test_cutedsl). + torch.testing.assert_close(recovered[:, NOPE_DIM:], ref[:, NOPE_DIM:]) diff --git a/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py b/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py index 9a5e478e31..a2085cd220 100644 --- a/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py +++ b/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py @@ -19,6 +19,7 @@ even/odd halves, producing (N_QUANT_BLOCKS, MXFP4_BLOCK/2) packed nibbles and N_QUANT_BLOCKS ue8m0 bytes. """ +from functools import lru_cache from typing import Any import torch @@ -296,6 +297,360 @@ def _fused_kv_compress_norm_rope_insert_sparse_attn( tl.store(bf16_ptr + rope_local, result.to(tl.bfloat16), mask=is_rope) +# ============================================================================= +# Split kernels variant of the head=512 compressor (deep cr=128 gather). +# - compress gather: instead of launching one program per token, split along +# the head dimension to maximize CU occupancy. The head dimension split +# does not require cross-group reduction +# - finalize norm rope quant store: same as the single pass kernel due to its +# per-token nature +# Mirrors the CUDA cutedsl split kernel where num_splits is occupancy-targeted. +# Currently only tested and validated on ROCm gfx950 +# ============================================================================= +@lru_cache(maxsize=1) +def _n_cu() -> int: + return torch.cuda.get_device_properties(0).multi_processor_count + + +def _pick_compress_num_splits( + num_actual: int, compress_ratio: int, head_dim: int +) -> int: + """Occupancy-targeted column splits for the cr>=128 head=512 compressor. + + Sizes the per-token fan-out so (estimated computing tokens) * num_splits ~ + #CU, capped by head tiling at a 32-wide min tile, as a power-of-2 divisor of + head_dim. + """ + max_splits = head_dim // 32 + est_compute = max(1, num_actual // compress_ratio) + target = -(-_n_cu() // est_compute) # ceil(#CU / est_compute) + ns = 1 + while ns * 2 <= min(target, max_splits) and head_dim % (ns * 2) == 0: + ns *= 2 + return ns + + +@triton.jit +def _compress_gather_split_sparse_attn( + state_cache_ptr, + state_cache_stride0, + state_cache_stride1, + positions_ptr, + slot_mapping_ptr, + token_to_req_indices_ptr, + block_table_ptr, + block_table_stride, + block_size, + scratch_ptr, + scratch_stride, + HEAD_SIZE: tl.constexpr, + STATE_WIDTH: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + NUM_SPLITS: tl.constexpr, + HEAD_TILE: tl.constexpr, # HEAD_SIZE // NUM_SPLITS +): + """Stage 1: per-(token, head-split) compress gather, write to fp32 scratch + + No-overlap gather (cr>=128) on rows [0, COMPRESS_RATIO) + """ + pid = tl.program_id(0) + token_idx = pid // NUM_SPLITS + split_idx = pid % NUM_SPLITS + + slot_id = tl.load(slot_mapping_ptr + token_idx) + if slot_id < 0: + return + position = tl.load(positions_ptr + token_idx) + if (position + 1) % COMPRESS_RATIO != 0: + return + req_idx = tl.load(token_to_req_indices_ptr + token_idx) + + start = position - COMPRESS_RATIO + 1 + rows = tl.arange(0, COMPRESS_RATIO) + pos = start + rows + mask_pos = pos >= 0 + block_numbers = tl.load( + block_table_ptr + req_idx * block_table_stride + pos // block_size, + mask=mask_pos, + other=0, + ).to(tl.int64) + block_offsets = pos % block_size + + col = split_idx * HEAD_TILE + tl.arange(0, HEAD_TILE) + row_base = ( + state_cache_ptr + + block_numbers * state_cache_stride0 + + block_offsets * state_cache_stride1 + ) + cmask = mask_pos[:, None] + + score = tl.load( + row_base[:, None] + STATE_WIDTH + col[None, :], + mask=cmask, + other=float("-inf"), + ) + score = tl.softmax(score, dim=0) + kv = tl.load(row_base[:, None] + col[None, :], mask=cmask, other=0.0) + compressed = tl.sum(kv * score, axis=0) # [HEAD_TILE] fp32 + tl.store(scratch_ptr + token_idx * scratch_stride + col, compressed) + + +@triton.jit +def _finalize_norm_rope_quant_store_sparse_attn( + scratch_ptr, + scratch_stride, + positions_ptr, + slot_mapping_ptr, + rms_norm_weight_ptr, + rms_norm_eps, + cos_sin_cache_ptr, + cos_sin_stride, + k_cache_ptr, + kv_slot_mapping_ptr, + kv_cache_block_size, + HEAD_SIZE: tl.constexpr, + TRITON_BLOCK_SIZE: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + ROPE_HEAD_DIM: tl.constexpr, + FP8_MAX: tl.constexpr, + QUANT_BLOCK: tl.constexpr, + TOKEN_STRIDE: tl.constexpr, + SCALE_DIM: tl.constexpr, + KV_BLOCK_STRIDE: tl.constexpr, +): + """Stage 2: read compressed_kv[512] from scratch buffer, then + RMSNorm + FP8 quant (nope) + RoPE + bf16 store + """ + token_idx = tl.program_id(0) + slot_id = tl.load(slot_mapping_ptr + token_idx) + if slot_id < 0: + return + position = tl.load(positions_ptr + token_idx) + if (position + 1) % COMPRESS_RATIO != 0: + return + + block = tl.arange(0, TRITON_BLOCK_SIZE) + mask = block < HEAD_SIZE + compressed_kv = tl.load( + scratch_ptr + token_idx * scratch_stride + block, mask=mask, other=0.0 + ) + + rms_w = tl.load(rms_norm_weight_ptr + block, mask=mask, other=0.0) + variance = tl.sum(compressed_kv * compressed_kv, axis=0) / HEAD_SIZE + rrms = tl.rsqrt(variance + rms_norm_eps) + normed = compressed_kv * rrms * rms_w + + kv_slot_idx = tl.load(kv_slot_mapping_ptr + token_idx) + if kv_slot_idx < 0: + return + kv_block_idx = kv_slot_idx // kv_cache_block_size + kv_pos_in_block = kv_slot_idx % kv_cache_block_size + cache_block_ptr = k_cache_ptr + kv_block_idx.to(tl.int64) * KV_BLOCK_STRIDE + fp8_ptr = cache_block_ptr + kv_pos_in_block * TOKEN_STRIDE + scale_ptr = ( + cache_block_ptr + + kv_cache_block_size * TOKEN_STRIDE + + kv_pos_in_block * SCALE_DIM + ) + + NOPE_HEAD_DIM: tl.constexpr = HEAD_SIZE - ROPE_HEAD_DIM + HALF_ROPE: tl.constexpr = ROPE_HEAD_DIM // 2 + N_QUANT_BLOCKS: tl.constexpr = TRITON_BLOCK_SIZE // QUANT_BLOCK + N_NOPE_BLOCKS: tl.constexpr = NOPE_HEAD_DIM // QUANT_BLOCK + INV_FP8_MAX: tl.constexpr = 1.0 / FP8_MAX + + quant_input = normed.to(tl.bfloat16).to(tl.float32) + quant_2d = tl.reshape(quant_input, (N_QUANT_BLOCKS, QUANT_BLOCK)) + block_absmax = tl.maximum(tl.max(tl.abs(quant_2d), axis=1), 1e-4) + raw_scales = block_absmax * INV_FP8_MAX + exponents = tl.ceil(tl.log2(raw_scales)) + inv_scales = tl.exp2(-exponents) + x_scaled = quant_2d * tl.reshape(inv_scales, (N_QUANT_BLOCKS, 1)) + x_clamped = tl.clamp(x_scaled, -FP8_MAX, FP8_MAX) + x_uint8 = tl.reshape( + x_clamped.to(tl.float8e4nv).to(tl.uint8, bitcast=True), + (TRITON_BLOCK_SIZE,), + ) + tl.store(fp8_ptr + block, x_uint8, mask=block < NOPE_HEAD_DIM) + + scale_idx = tl.arange(0, N_QUANT_BLOCKS) + encoded = tl.maximum(tl.minimum(exponents + 127.0, 255.0), 0.0) + tl.store( + scale_ptr + scale_idx, encoded.to(tl.uint8), mask=scale_idx < N_NOPE_BLOCKS + ) + tl.store(scale_ptr + N_NOPE_BLOCKS, tl.zeros((), dtype=tl.uint8)) + + NUM_PAIRS: tl.constexpr = TRITON_BLOCK_SIZE // 2 + NOPE_PAIRS: tl.constexpr = NOPE_HEAD_DIM // 2 + even, odd = tl.split(tl.reshape(normed, (NUM_PAIRS, 2))) + pair_idx = tl.arange(0, NUM_PAIRS) + rope_pair_local = pair_idx - NOPE_PAIRS + is_rope_pair = rope_pair_local >= 0 + cs_idx = tl.maximum(rope_pair_local, 0) + compressed_pos = (position // COMPRESS_RATIO) * COMPRESS_RATIO + cache_base = cos_sin_cache_ptr + compressed_pos * cos_sin_stride + cos_v = tl.load(cache_base + cs_idx, mask=is_rope_pair, other=1.0) + sin_v = tl.load(cache_base + HALF_ROPE + cs_idx, mask=is_rope_pair, other=0.0) + new_even = even * cos_v - odd * sin_v + new_odd = odd * cos_v + even * sin_v + result = tl.interleave(new_even, new_odd) + bf16_ptr = (fp8_ptr + NOPE_HEAD_DIM).to(tl.pointer_type(tl.bfloat16)) + rope_local = block - NOPE_HEAD_DIM + is_rope = (block >= NOPE_HEAD_DIM) & mask + tl.store(bf16_ptr + rope_local, result.to(tl.bfloat16), mask=is_rope) + + +def _launch_two_stage_sparse_attn_compressor( + state_cache: torch.Tensor, + token_to_req_indices: torch.Tensor, + positions: torch.Tensor, + slot_mapping: torch.Tensor, + block_table: torch.Tensor, + block_size: int, + state_width: int, + compress_ratio: int, + cos_sin_cache: torch.Tensor, + kv_cache: torch.Tensor, + kv_slot_mapping: torch.Tensor, + rms_norm_weight: torch.Tensor, + rms_norm_eps: float, + quant_block: int, + token_stride: int, + scale_dim: int, + head_dim: int, + rope_head_dim: int, + num_actual: int, + compress_scratch: torch.Tensor, +) -> None: + num_splits = _pick_compress_num_splits(num_actual, compress_ratio, head_dim) + head_tile = head_dim // num_splits + scratch = compress_scratch[:num_actual] + _compress_gather_split_sparse_attn[(num_actual * num_splits,)]( + state_cache, + state_cache.stride(0), + state_cache.stride(1), + positions, + slot_mapping, + token_to_req_indices, + block_table, + block_table.stride(0), + block_size, + scratch, + scratch.stride(0), + HEAD_SIZE=head_dim, + STATE_WIDTH=state_width, + COMPRESS_RATIO=compress_ratio, + NUM_SPLITS=num_splits, + HEAD_TILE=head_tile, + ) + _finalize_norm_rope_quant_store_sparse_attn[(num_actual,)]( + scratch, + scratch.stride(0), + positions, + slot_mapping, + rms_norm_weight, + rms_norm_eps, + cos_sin_cache, + cos_sin_cache.stride(0), + kv_cache, + kv_slot_mapping, + kv_cache.shape[1], + HEAD_SIZE=head_dim, + TRITON_BLOCK_SIZE=triton.next_power_of_2(head_dim), + COMPRESS_RATIO=compress_ratio, + ROPE_HEAD_DIM=rope_head_dim, + FP8_MAX=448.0, + QUANT_BLOCK=quant_block, + TOKEN_STRIDE=token_stride, + SCALE_DIM=scale_dim, + KV_BLOCK_STRIDE=kv_cache.stride(0), + ) + + +def compress_norm_rope_store_two_stage_triton( + state_cache: torch.Tensor, + num_actual: int, + token_to_req_indices: torch.Tensor, + positions: torch.Tensor, + slot_mapping: torch.Tensor, + block_table: torch.Tensor, + block_size: int, + state_width: int, + cos_sin_cache: torch.Tensor, + kv_cache: torch.Tensor, + k_cache_metadata: Any, + pdl_kwargs: dict, + head_dim: int, + rope_head_dim: int, + compress_ratio: int, + overlap: bool, + use_fp4_cache: bool, + rms_norm_weight: torch.Tensor, + rms_norm_eps: float, + quant_block: int, + token_stride: int, + scale_dim: int, + num_decode_tokens: int, + compress_scratch: torch.Tensor, +) -> None: + """Two-stage split compressor dispatch for head=512 cr>=128 (no-overlap) + + Run the occupancy-fanned two-stage split for prefill [num_decodee_tokens:] + to fill the CUs, and use the original single-pass launcher + for decode [0, num_decode_tokens) + """ + num_decodes = min(max(num_decode_tokens, 0), num_actual) + num_prefills = num_actual - num_decodes + if num_prefills > 0: + _launch_two_stage_sparse_attn_compressor( + state_cache=state_cache, + token_to_req_indices=token_to_req_indices[num_decodes:], + positions=positions[num_decodes:], + slot_mapping=slot_mapping[num_decodes:], + block_table=block_table, + block_size=block_size, + state_width=state_width, + compress_ratio=compress_ratio, + cos_sin_cache=cos_sin_cache, + kv_cache=kv_cache, + kv_slot_mapping=k_cache_metadata.slot_mapping[num_decodes:], + rms_norm_weight=rms_norm_weight, + rms_norm_eps=rms_norm_eps, + quant_block=quant_block, + token_stride=token_stride, + scale_dim=scale_dim, + head_dim=head_dim, + rope_head_dim=rope_head_dim, + num_actual=num_prefills, + compress_scratch=compress_scratch, + ) + if num_decodes > 0: + compress_norm_rope_store_triton( + state_cache=state_cache, + num_actual=num_decodes, + token_to_req_indices=token_to_req_indices, + positions=positions, + slot_mapping=slot_mapping, + block_table=block_table, + block_size=block_size, + state_width=state_width, + cos_sin_cache=cos_sin_cache, + kv_cache=kv_cache, + k_cache_metadata=k_cache_metadata, + pdl_kwargs=pdl_kwargs, + head_dim=head_dim, + rope_head_dim=rope_head_dim, + compress_ratio=compress_ratio, + overlap=overlap, + use_fp4_cache=use_fp4_cache, + rms_norm_weight=rms_norm_weight, + rms_norm_eps=rms_norm_eps, + quant_block=quant_block, + token_stride=token_stride, + scale_dim=scale_dim, + ) + + # ============================================================================= # Indexer path (head=128, all FP8, single quant block) # ============================================================================= diff --git a/vllm/models/deepseek_v4/compressor.py b/vllm/models/deepseek_v4/compressor.py index 24838c237c..13f327f6bc 100644 --- a/vllm/models/deepseek_v4/compressor.py +++ b/vllm/models/deepseek_v4/compressor.py @@ -14,6 +14,7 @@ from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import MergedColumnParallelLinear from vllm.models.deepseek_v4.common.ops.fused_compress_quant_cache import ( compress_norm_rope_store_triton, + compress_norm_rope_store_two_stage_triton, ) from vllm.models.deepseek_v4.common.ops.fused_indexer_q import MXFP4_BLOCK_SIZE from vllm.models.deepseek_v4.common.ops.save_partial_states import ( @@ -27,6 +28,7 @@ from vllm.v1.attention.backend import ( CommonAttentionMetadata, MultipleOf, ) +from vllm.v1.attention.backends.utils import split_decodes_and_prefills from vllm.v1.kv_cache_interface import ( KVCacheSpec, MLAAttentionSpec, @@ -34,6 +36,12 @@ from vllm.v1.kv_cache_interface import ( ) +def _prefer_two_stage_compressor() -> bool: + # Platforms that favor the triton variant of two-stage compressor split. + # Currently only tested on ROCm + return current_platform.is_rocm() + + class CompressorBackend(AttentionBackend): def __init__(self): super().__init__() @@ -81,6 +89,7 @@ class CompressorMetadata: block_size: int token_to_req_indices: torch.Tensor | None = None # [num_tokens] + num_decode_tokens: int | None = None class CompressorMetadataBuilder(AttentionMetadataBuilder): @@ -107,11 +116,17 @@ class CompressorMetadataBuilder(AttentionMetadataBuilder): token_to_req_indices = common_attn_metadata.token_to_req_indices( self.token_to_req_indices ) + num_decode_tokens = None + if _prefer_two_stage_compressor(): + _, _, num_decode_tokens, _ = split_decodes_and_prefills( + common_attn_metadata, decode_threshold=1 + ) return CompressorMetadata( block_table=common_attn_metadata.block_table_tensor.clamp_(min=0), slot_mapping=common_attn_metadata.slot_mapping, block_size=self.block_size, token_to_req_indices=token_to_req_indices, + num_decode_tokens=num_decode_tokens, ) @@ -213,6 +228,25 @@ class DeepseekCompressor(nn.Module): self.overlap = compress_ratio == 4 self.coff = 1 + self.overlap + # The head=512 cr>=128 no-overlap deep gather uses the two-stage + # compressor, which needs an fp32 scratch [max_batched, 512] for + # the intermediate compressed_kv. + # Currently only tested on ROCm + self._use_two_stage_fused_compressor = ( + _prefer_two_stage_compressor() and head_dim == 512 and not self.overlap + ) + self.max_num_batched_tokens = ( + vllm_config.scheduler_config.max_num_batched_tokens + ) + self._compress_scratch: torch.Tensor | None = None + if self._use_two_stage_fused_compressor: + self._compress_scratch = torch.empty( + self.max_num_batched_tokens, + self.head_dim, + dtype=torch.float32, + device=self.device, + ) + state_dtype = torch.float32 self.ape = nn.Parameter( torch.empty( @@ -364,6 +398,15 @@ class DeepseekCompressor(nn.Module): store_full_fp8=store_full_fp8, fp8_scale=fp8_scale, ) + elif self._use_two_stage_fused_compressor: + # head=512 cr>=128 (no overlap): two-pass split compressor on the + # prefill suffix, single-pass on the decode prefix. + assert state_metadata.num_decode_tokens is not None + compress_norm_rope_store_fn = compress_norm_rope_store_two_stage_triton + extra_kwargs = { + "num_decode_tokens": state_metadata.num_decode_tokens, + "compress_scratch": self._compress_scratch, + } else: # Indexer path (head_dim == 128) or non-CUDA GPUs (AMD, XPU, etc.). compress_norm_rope_store_fn = compress_norm_rope_store_triton