iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: kliuae <kuanfu.liu@embeddedllm.com>
This commit is contained in:
kliuae
2026-07-15 23:31:51 +00:00
committed by GitHub
parent 2bd8957627
commit eb33ff34dd
3 changed files with 528 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/test_compressor_kv_cache.py
+130
View File
@@ -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:])
vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py
+355
View File
@@ -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)
# =============================================================================
vllm/models/deepseek_v4/compressor.py
+43
View File
@@ -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