iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity
Files
382970ee6ca490aeaaaf4e32c53695b581ff61ba
vLLM/tests/kernels/test_deepseek_v4_cpu_kernels.py
T

1495 lines
53 KiB
Python
Raw Blame History

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""CPU-kernel-vs-eager-reference parity tests for the DeepSeek-V4 CPU port's
sgl-kernels ops (csrc/cpu/sgl-kernels/). Each test carries its own reference
reimplementation of the kernel's math rather than depending on production
model code.
"""
from unittest import mock
import pytest
import torch
import vllm.model_executor.kernels.mhc as mhc_kernels
from tests.kernels.test_fused_deepseek_v4_qnorm_rope_kv_insert import (
HEAD_DIM,
NOPE_DIM,
QUANT_BLOCK,
ROPE_DIM,
make_cos_sin_cache,
)
from tests.kernels.utils import bf16_ulp_distance, fp8_ulp_distance
from vllm._custom_ops import (
compress_norm_rope_store_cpu,
compress_norm_rope_store_indexer_cpu,
flash_mla_with_kvcache_cpu,
fp8_paged_mqa_logits_cpu,
fused_indexer_q_rope_quant_cpu,
fused_qnorm_rope_kv_insert_cpu,
inverse_gptj_rope_o_proj_cpu,
save_partial_states_cpu,
)
from vllm.models.deepseek_v4.common.ops import fused_indexer_q_rope_quant
from vllm.platforms import current_platform
from vllm.utils.torch_utils import set_random_seed
DEVICE = current_platform.device_type
# Main attention cache (fp8_ds_mla): NoPE fp8 in QUANT_BLOCK-wide UE8M0 blocks,
# RoPE bf16, one scale byte per block (+1 pad byte) per token.
SCALE_BYTES_PER_TOKEN = NOPE_DIM // QUANT_BLOCK + 1
TOKEN_DATA_BYTES = NOPE_DIM + ROPE_DIM * 2
FP8_MAX = 448.0
# Indexer K-cache: all 128 dims in one fp8 block, one raw fp32 scale per token.
INDEXER_HEAD_DIM = 128
INDEXER_NOPE_DIM = INDEXER_HEAD_DIM - ROPE_DIM
INDEXER_SCALE_BYTES = 4
def _op_available(*names: str) -> bool:
return all(hasattr(torch.ops._C, name) for name in names)
pytestmark = pytest.mark.skipif(
not current_platform.is_cpu()
or not _op_available(
"save_partial_states_cpu",
"compress_norm_rope_store_cpu",
"compress_norm_rope_store_indexer_cpu",
"flash_mla_with_kvcache_cpu",
"fused_qnorm_rope_kv_insert_cpu",
"inverse_gptj_rope_o_proj_cpu",
"fp8_paged_mqa_logits_cpu",
"fused_indexer_q_rope_quant_cpu",
),
reason="CPU not available or DeepSeek-V4 CPU kernels not built in",
)
def _apply_gptj_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
x_even = x[..., 0::2]
x_odd = x[..., 1::2]
out_even = x_even * cos - x_odd * sin
out_odd = x_even * sin + x_odd * cos
return torch.stack((out_even, out_odd), dim=-1).flatten(-2)
def _make_cache(num_blocks: int, block_size: int) -> torch.Tensor:
block_bytes = block_size * TOKEN_DATA_BYTES + block_size * SCALE_BYTES_PER_TOKEN
return torch.zeros(num_blocks, block_bytes, dtype=torch.uint8)
def _decode_token_regions(
cache: torch.Tensor, slot_mapping: torch.Tensor, block_size: int
):
"""Slice out, per valid token, the raw NoPE-fp8 / RoPE-bf16 / scale byte
regions this layout defines."""
valid = slot_mapping >= 0
idx = valid.nonzero(as_tuple=True)[0]
slots = slot_mapping[idx]
block_idx = torch.div(slots, block_size, rounding_mode="floor")
pos_in_block = slots % block_size
nope, rope, scale = [], [], []
for b, p in zip(block_idx.tolist(), pos_in_block.tolist()):
token_off = p * TOKEN_DATA_BYTES
nope.append(cache[b, token_off : token_off + NOPE_DIM])
rope.append(cache[b, token_off + NOPE_DIM : token_off + TOKEN_DATA_BYTES])
scale_off = block_size * TOKEN_DATA_BYTES + p * SCALE_BYTES_PER_TOKEN
scale.append(cache[b, scale_off : scale_off + SCALE_BYTES_PER_TOKEN - 1])
if not nope:
empty = lambda n: torch.empty(0, n, dtype=torch.uint8)
return empty(NOPE_DIM), empty(ROPE_DIM * 2), empty(SCALE_BYTES_PER_TOKEN - 1)
return torch.stack(nope), torch.stack(rope), torch.stack(scale)
def _assert_cache_parity(cache_fused, cache_ref, slot_mapping, block_size):
"""NoPE fp8 + the UE8M0 scale byte are a deterministic quant of an
un-rotated bf16 value, so they match tightly; the RoPE region is rotated
in fp32 before its bf16 store and eager/kernel can round a tie to
opposite sides, so allow <=1 ULP there."""
nope_ref, rope_ref, scale_ref = _decode_token_regions(
cache_ref, slot_mapping, block_size
)
nope_fused, rope_fused, scale_fused = _decode_token_regions(
cache_fused, slot_mapping, block_size
)
if nope_ref.numel() == 0:
return
max_ulp = int(fp8_ulp_distance(nope_fused, nope_ref).max().item())
assert max_ulp <= 1, f"NoPE fp8 differs by {max_ulp} ULP (>1)"
rope_ref_bf16 = rope_ref.contiguous().view(torch.bfloat16)
rope_fused_bf16 = rope_fused.contiguous().view(torch.bfloat16)
max_ulp_rope = int(bf16_ulp_distance(rope_fused_bf16, rope_ref_bf16).max().item())
assert max_ulp_rope <= 1, f"RoPE bf16 differs by {max_ulp_rope} ULP (>1)"
scale_diff = (scale_fused.int() - scale_ref.int()).abs().max().item()
assert scale_diff <= 1, f"scale byte differs by {scale_diff} (>1)"
def _make_indexer_cache(num_blocks: int, block_size: int) -> torch.Tensor:
block_bytes = block_size * INDEXER_HEAD_DIM + block_size * INDEXER_SCALE_BYTES
return torch.zeros(num_blocks, block_bytes, dtype=torch.uint8)
def _assert_indexer_cache_parity(cache_fused, cache_ref, slot_mapping, block_size):
valid = slot_mapping >= 0
idx = valid.nonzero(as_tuple=True)[0]
if idx.numel() == 0:
return
slots = slot_mapping[idx]
block_idx = torch.div(slots, block_size, rounding_mode="floor")
pos_in_block = slots % block_size
fp8_ref, scale_ref = [], []
fp8_fused, scale_fused = [], []
for b, p in zip(block_idx.tolist(), pos_in_block.tolist()):
token_off = p * INDEXER_HEAD_DIM
scale_off = block_size * INDEXER_HEAD_DIM + p * INDEXER_SCALE_BYTES
fp8_ref.append(cache_ref[b, token_off : token_off + INDEXER_HEAD_DIM])
fp8_fused.append(cache_fused[b, token_off : token_off + INDEXER_HEAD_DIM])
scale_ref.append(cache_ref[b, scale_off : scale_off + INDEXER_SCALE_BYTES])
scale_fused.append(cache_fused[b, scale_off : scale_off + INDEXER_SCALE_BYTES])
fp8_ref_t = torch.stack(fp8_ref).view(torch.float8_e4m3fn)
fp8_fused_t = torch.stack(fp8_fused).view(torch.float8_e4m3fn)
torch.testing.assert_close(fp8_fused_t.float(), fp8_ref_t.float())
scale_ref_t = torch.stack(scale_ref).view(torch.float32)
scale_fused_t = torch.stack(scale_fused).view(torch.float32)
torch.testing.assert_close(scale_fused_t, scale_ref_t)
# ---------------------------------------------------------------------------
# save_partial_states_cpu / compress_norm_rope_store_cpu (compressor.cpp)
# ---------------------------------------------------------------------------
#
# The compressor's state cache is pooled with the main KV cache's physical
# pages and page-aligned, so it is genuinely 3D with a block-to-block stride
# that need not equal block_size * width -- built here via torch.as_strided
# over an over-sized flat buffer rather than a plain contiguous tensor, so a
# regression back to flattening via .view() would be caught.
def _make_padded_block_cache(
num_blocks: int, block_size: int, width: int, pad_elems: int
) -> torch.Tensor:
page_stride = block_size * width + pad_elems
flat = torch.randn(num_blocks, page_stride, dtype=torch.float32)
return torch.as_strided(
flat, (num_blocks, block_size, width), (page_stride, width, 1)
)
def _dup_padded_block_cache(
num_blocks: int, block_size: int, width: int, pad_elems: int
) -> tuple[torch.Tensor, torch.Tensor]:
"""Two independently-stored padded caches with identical values, for ops
that mutate the cache in place."""
page_stride = block_size * width + pad_elems
template = torch.randn(num_blocks, page_stride, dtype=torch.float32)
def _view(flat: torch.Tensor) -> torch.Tensor:
return torch.as_strided(
flat, (num_blocks, block_size, width), (page_stride, width, 1)
)
return _view(template.clone()), _view(template.clone())
def _run_save_partial_states_eager(
state_cache, kv, score, ape, positions, slot_mapping, compress_ratio
):
"""Reference: state_cache[block_idx, pos_in_block, :W] = kv;
state_cache[..., W:] = score + ape[position % compress_ratio]; skip
slot < 0. Mutates state_cache in place, matching the op's contract."""
state_width = kv.shape[-1]
block_size = state_cache.shape[1]
valid = slot_mapping >= 0
idx = valid.nonzero(as_tuple=True)[0]
slots = slot_mapping[idx]
block_idx = slots // block_size
pos_in_block = slots % block_size
state_cache[block_idx, pos_in_block, :state_width] = kv[idx]
ape_row = ape[positions[idx] % compress_ratio]
state_cache[block_idx, pos_in_block, state_width:] = score[idx] + ape_row
@pytest.mark.parametrize("pad_elems", [0, 24])
@pytest.mark.parametrize("num_tokens", [1, 17])
def test_save_partial_states_cpu_matches_eager(num_tokens: int, pad_elems: int):
torch.manual_seed(0)
compress_ratio = 4
block_size = 4
state_width = 512
width = 2 * state_width
num_blocks = (num_tokens + 4 + block_size - 1) // block_size + 1
num_slots = num_blocks * block_size
kv = torch.randn(num_tokens, state_width, dtype=torch.float32)
score = torch.randn(num_tokens, state_width, dtype=torch.float32)
ape = torch.randn(compress_ratio, state_width, dtype=torch.float32)
positions = torch.randint(0, 4096, (num_tokens,), dtype=torch.int64)
ref_cache, fused_cache = _dup_padded_block_cache(
num_blocks, block_size, width, pad_elems
)
slot_mapping = torch.randperm(num_slots)[:num_tokens].to(torch.int64)
if num_tokens > 1:
slot_mapping[-1] = -1 # exercise the "skip invalid slot" branch
_run_save_partial_states_eager(
ref_cache, kv, score, ape, positions, slot_mapping, compress_ratio
)
save_partial_states_cpu(kv, score, ape, positions, fused_cache, slot_mapping)
torch.testing.assert_close(fused_cache, ref_cache)
def _run_compress_eager(
state_cache,
gather_slots,
positions,
kv_slot_mapping,
rms_norm_weight,
rms_norm_eps,
cos_sin_cache,
kv_cache,
kv_cache_block_size,
compress_ratio,
):
"""Reference for the per-channel-softmax-pooling + RMSNorm + NoPE-UE8M0
quant + GPT-J RoPE math writing the main (head_dim=512) fp8_ds_mla
layout."""
num_tokens, window = gather_slots.shape
state_width = state_cache.shape[-1] // 2
state_block_size = state_cache.shape[1]
num_quant_blocks = NOPE_DIM // QUANT_BLOCK
w_idx = torch.arange(window)
head_offset = torch.where(w_idx >= compress_ratio, HEAD_DIM, 0)
for t in range(num_tokens):
position = int(positions[t])
if (position + 1) % compress_ratio != 0:
continue
kv_slot = int(kv_slot_mapping[t])
if kv_slot < 0:
continue
slots = gather_slots[t]
valid_w = slots >= 0
safe_slots = slots.clamp(min=0)
block_idx = safe_slots // state_block_size
pos_in_block = safe_slots % state_block_size
rows = state_cache[block_idx, pos_in_block] # [window, 2*state_width]
col = head_offset.unsqueeze(-1) + torch.arange(HEAD_DIM)
kv_vals = torch.gather(rows, 1, col)
score_col = state_width + col
score_vals = torch.gather(rows, 1, score_col)
score_vals = torch.where(
valid_w.unsqueeze(-1), score_vals, torch.full_like(score_vals, -torch.inf)
)
weights = torch.softmax(score_vals, dim=0)
compressed = (kv_vals * weights).sum(dim=0) # [HEAD_DIM]
variance = compressed.pow(2).mean()
normed = compressed * torch.rsqrt(variance + rms_norm_eps) * rms_norm_weight
quant_input = normed.to(torch.bfloat16).to(torch.float32)
compressed_pos = (position // compress_ratio) * compress_ratio
cos = cos_sin_cache[compressed_pos, : ROPE_DIM // 2]
sin = cos_sin_cache[compressed_pos, ROPE_DIM // 2 :]
rotated_rope = _apply_gptj_rope(normed[NOPE_DIM:].clone(), cos, sin)
nope_vals = quant_input[:NOPE_DIM].view(num_quant_blocks, QUANT_BLOCK)
absmax = nope_vals.abs().amax(dim=-1).clamp(min=1e-4)
exponent = torch.ceil(torch.log2(absmax / FP8_MAX))
inv_scale = torch.exp2(-exponent)
scaled = (nope_vals * inv_scale.unsqueeze(-1)).clamp(-FP8_MAX, FP8_MAX)
nope_fp8_bytes = scaled.to(torch.float8_e4m3fn).reshape(-1).view(torch.uint8)
scale_bytes = (exponent + 127.0).clamp(0, 255).to(torch.uint8)
rope_bf16_bytes = rotated_rope.to(torch.bfloat16).contiguous().view(torch.uint8)
block_idx_kv = kv_slot // kv_cache_block_size
pos_in_block_kv = kv_slot % kv_cache_block_size
token_off = pos_in_block_kv * TOKEN_DATA_BYTES
kv_cache[block_idx_kv, token_off : token_off + NOPE_DIM] = nope_fp8_bytes
kv_cache[block_idx_kv, token_off + NOPE_DIM : token_off + TOKEN_DATA_BYTES] = (
rope_bf16_bytes
)
scale_off = (
kv_cache_block_size * TOKEN_DATA_BYTES
+ pos_in_block_kv * SCALE_BYTES_PER_TOKEN
)
kv_cache[block_idx_kv, scale_off : scale_off + num_quant_blocks] = scale_bytes
kv_cache[block_idx_kv, scale_off + num_quant_blocks] = 0
def _make_compress_inputs(
num_tokens: int,
compress_ratio: int,
overlap: bool,
pad_elems: int = 24,
max_pos: int = 4096,
invalid_window_frac: float = 0.1,
head_dim: int = HEAD_DIM,
):
torch.manual_seed(0)
coff = 1 + int(overlap)
state_width = coff * head_dim
window = coff * compress_ratio
state_block_size = 4 if compress_ratio == 4 else 8
num_state_slots = num_tokens * window + 8
num_state_blocks = (num_state_slots + state_block_size - 1) // state_block_size + 1
state_cache = _make_padded_block_cache(
num_state_blocks, state_block_size, 2 * state_width, pad_elems
)
gather_slots = torch.randint(
0, num_state_blocks * state_block_size, (num_tokens, window)
).to(torch.int64)
invalid_mask = torch.rand(num_tokens, window) < invalid_window_frac
# Never invalidate every column of a row (would force max=-inf -> NaN).
invalid_mask[:, -1] = False
gather_slots = torch.where(
invalid_mask, torch.full_like(gather_slots, -1), gather_slots
)
base = (torch.arange(num_tokens) + 1) * compress_ratio
is_boundary = torch.arange(num_tokens) % 2 == 0
positions = torch.where(is_boundary, base - 1, base + 1).to(torch.int64)
kv_slot_mapping = torch.arange(num_tokens, dtype=torch.int64)
if num_tokens > 1:
kv_slot_mapping[1] = -1 # exercise the "skip invalid output slot" branch
rms_norm_weight = torch.randn(head_dim, dtype=torch.float32)
cos_sin_cache = make_cos_sin_cache(max_pos, ROPE_DIM, torch.float32, "cpu")
return (
state_cache,
gather_slots,
positions,
kv_slot_mapping,
rms_norm_weight,
cos_sin_cache,
is_boundary,
)
@pytest.mark.parametrize("compress_ratio,overlap", [(128, False), (4, True)])
@pytest.mark.parametrize("num_tokens", [1, 17])
def test_compress_norm_rope_store_cpu_matches_eager(
compress_ratio: int, overlap: bool, num_tokens: int
):
eps = 1e-6
block_size = 16
(
state_cache,
gather_slots,
positions,
kv_slot_mapping,
rms_norm_weight,
cos_sin_cache,
is_boundary,
) = _make_compress_inputs(num_tokens, compress_ratio, overlap)
num_blocks = (num_tokens + block_size - 1) // block_size + 1
kv_cache_ref = _make_cache(num_blocks, block_size)
kv_cache_fused = kv_cache_ref.clone()
_run_compress_eager(
state_cache,
gather_slots,
positions,
kv_slot_mapping,
rms_norm_weight,
eps,
cos_sin_cache,
kv_cache_ref,
block_size,
compress_ratio,
)
compress_norm_rope_store_cpu(
state_cache,
gather_slots,
positions,
kv_slot_mapping,
rms_norm_weight,
eps,
cos_sin_cache,
kv_cache_fused,
block_size,
compress_ratio,
)
# Only tokens that hit the compression boundary (and have a valid output
# slot) were written -- restrict parity checks to those.
written_slot_mapping = torch.where(
is_boundary, kv_slot_mapping, torch.full_like(kv_slot_mapping, -1)
)
_assert_cache_parity(kv_cache_fused, kv_cache_ref, written_slot_mapping, block_size)
def test_compress_norm_rope_store_cpu_non_contiguous_kv_cache_row_stride():
"""The main (output) kv_cache is a separate physical buffer from the
state cache and can independently be a strided view."""
eps = 1e-6
compress_ratio, overlap = 128, False
num_tokens = 5
block_size = 16
cache_pad = 32
(
state_cache,
gather_slots,
positions,
kv_slot_mapping,
rms_norm_weight,
cos_sin_cache,
is_boundary,
) = _make_compress_inputs(num_tokens, compress_ratio, overlap)
num_blocks = (num_tokens + block_size - 1) // block_size + 1
kv_cache_ref = _make_cache(num_blocks, block_size)
block_bytes = kv_cache_ref.shape[1]
padded_cache_ref = torch.zeros(
num_blocks, block_bytes + cache_pad, dtype=torch.uint8
)
kv_cache_ref_view = padded_cache_ref[:, :block_bytes]
kv_cache_ref_view.copy_(kv_cache_ref)
padded_cache_fused = padded_cache_ref.clone()
kv_cache_fused_view = padded_cache_fused[:, :block_bytes]
_run_compress_eager(
state_cache,
gather_slots,
positions,
kv_slot_mapping,
rms_norm_weight,
eps,
cos_sin_cache,
kv_cache_ref_view,
block_size,
compress_ratio,
)
compress_norm_rope_store_cpu(
state_cache,
gather_slots,
positions,
kv_slot_mapping,
rms_norm_weight,
eps,
cos_sin_cache,
kv_cache_fused_view,
block_size,
compress_ratio,
)
written_slot_mapping = torch.where(
is_boundary, kv_slot_mapping, torch.full_like(kv_slot_mapping, -1)
)
_assert_cache_parity(
kv_cache_fused_view, kv_cache_ref_view, written_slot_mapping, block_size
)
def _run_compress_indexer_eager(
state_cache,
gather_slots,
positions,
kv_slot_mapping,
rms_norm_weight,
rms_norm_eps,
cos_sin_cache,
kv_cache,
kv_cache_block_size,
compress_ratio,
):
"""Same softmax-pool + RMSNorm + GPT-J RoPE prologue as the head=512
path, but the indexer's cache quantizes all 128 post-RoPE dims as one
FP8 block with a single raw fp32 scale per token (matching
_indexer_k_quant_and_cache_cpu's layout), not 7 separate UE8M0 blocks."""
num_tokens, window = gather_slots.shape
state_width = state_cache.shape[-1] // 2
state_block_size = state_cache.shape[1]
w_idx = torch.arange(window)
head_offset = torch.where(w_idx >= compress_ratio, INDEXER_HEAD_DIM, 0)
for t in range(num_tokens):
position = int(positions[t])
if (position + 1) % compress_ratio != 0:
continue
kv_slot = int(kv_slot_mapping[t])
if kv_slot < 0:
continue
slots = gather_slots[t]
valid_w = slots >= 0
safe_slots = slots.clamp(min=0)
block_idx = safe_slots // state_block_size
pos_in_block = safe_slots % state_block_size
rows = state_cache[block_idx, pos_in_block] # [window, 2*state_width]
col = head_offset.unsqueeze(-1) + torch.arange(INDEXER_HEAD_DIM)
kv_vals = torch.gather(rows, 1, col)
score_col = state_width + col
score_vals = torch.gather(rows, 1, score_col)
score_vals = torch.where(
valid_w.unsqueeze(-1), score_vals, torch.full_like(score_vals, -torch.inf)
)
weights = torch.softmax(score_vals, dim=0)
compressed = (kv_vals * weights).sum(dim=0) # [INDEXER_HEAD_DIM]
variance = compressed.pow(2).mean()
normed = compressed * torch.rsqrt(variance + rms_norm_eps) * rms_norm_weight
compressed_pos = (position // compress_ratio) * compress_ratio
cos = cos_sin_cache[compressed_pos, : ROPE_DIM // 2]
sin = cos_sin_cache[compressed_pos, ROPE_DIM // 2 :]
rotated_rope = _apply_gptj_rope(normed[INDEXER_NOPE_DIM:].clone(), cos, sin)
full = torch.cat((normed[:INDEXER_NOPE_DIM], rotated_rope))
quant_input = full.to(torch.bfloat16).to(torch.float32)
absmax = quant_input.abs().amax().clamp(min=1e-4)
exponent = torch.ceil(torch.log2(absmax / FP8_MAX))
scale = torch.exp2(exponent)
scaled = (quant_input / scale).clamp(-FP8_MAX, FP8_MAX)
fp8_bytes = scaled.to(torch.float8_e4m3fn).view(torch.uint8)
block_idx_kv = kv_slot // kv_cache_block_size
pos_in_block_kv = kv_slot % kv_cache_block_size
token_off = pos_in_block_kv * INDEXER_HEAD_DIM
kv_cache[block_idx_kv, token_off : token_off + INDEXER_HEAD_DIM] = fp8_bytes
scale_off = (
kv_cache_block_size * INDEXER_HEAD_DIM
+ pos_in_block_kv * INDEXER_SCALE_BYTES
)
kv_cache[block_idx_kv, scale_off : scale_off + INDEXER_SCALE_BYTES] = (
scale.reshape(1).view(torch.uint8)
)
@pytest.mark.parametrize("compress_ratio,overlap", [(128, False), (4, True)])
@pytest.mark.parametrize("num_tokens", [1, 17])
def test_compress_norm_rope_store_indexer_cpu_matches_eager(
compress_ratio: int, overlap: bool, num_tokens: int
):
eps = 1e-6
block_size = 16
(
state_cache,
gather_slots,
positions,
kv_slot_mapping,
rms_norm_weight,
cos_sin_cache,
is_boundary,
) = _make_compress_inputs(
num_tokens, compress_ratio, overlap, head_dim=INDEXER_HEAD_DIM
)
num_blocks = (num_tokens + block_size - 1) // block_size + 1
kv_cache_ref = _make_indexer_cache(num_blocks, block_size)
kv_cache_fused = kv_cache_ref.clone()
_run_compress_indexer_eager(
state_cache,
gather_slots,
positions,
kv_slot_mapping,
rms_norm_weight,
eps,
cos_sin_cache,
kv_cache_ref,
block_size,
compress_ratio,
)
compress_norm_rope_store_indexer_cpu(
state_cache,
gather_slots,
positions,
kv_slot_mapping,
rms_norm_weight,
eps,
cos_sin_cache,
kv_cache_fused,
block_size,
compress_ratio,
)
written_slot_mapping = torch.where(
is_boundary, kv_slot_mapping, torch.full_like(kv_slot_mapping, -1)
)
_assert_indexer_cache_parity(
kv_cache_fused, kv_cache_ref, written_slot_mapping, block_size
)
# ---------------------------------------------------------------------------
# flash_mla_with_kvcache_cpu (flash_mla.cpp): fused sparse MQA decode attn
# ---------------------------------------------------------------------------
def _pack_flash_mla_cache(kv_f32: torch.Tensor, block_size: int):
"""Quantize kv_f32 ([N, HEAD_DIM]) into the fp8_ds_mla byte layout and
return (cache, dequant_ref) -- the exact post-quantization value the
kernel's dequant should reproduce. Slot id == row index, packed
contiguously from block 0."""
n = kv_f32.shape[0]
num_blocks = max(1, -(-n // block_size))
cache = _make_cache(num_blocks, block_size)
if n == 0:
return cache, kv_f32.new_zeros(0, HEAD_DIM)
num_quant_blocks = NOPE_DIM // QUANT_BLOCK
nope = kv_f32[:, :NOPE_DIM].reshape(n, num_quant_blocks, QUANT_BLOCK)
absmax = nope.abs().amax(dim=-1).clamp(min=1e-4)
exponent = torch.ceil(torch.log2(absmax / FP8_MAX))
inv_scale = torch.exp2(-exponent)
scaled = (nope * inv_scale.unsqueeze(-1)).clamp(-FP8_MAX, FP8_MAX)
nope_fp8 = scaled.to(torch.float8_e4m3fn)
scale = torch.exp2(exponent)
nope_dequant = (nope_fp8.float() * scale.unsqueeze(-1)).reshape(n, NOPE_DIM)
rope_bf16 = kv_f32[:, NOPE_DIM:].to(torch.bfloat16)
dequant_ref = torch.cat([nope_dequant, rope_bf16.float()], dim=-1)
slot_ids = torch.arange(n)
block_idx = slot_ids // block_size
pos_in_block = slot_ids % block_size
nope_fp8_bytes = nope_fp8.reshape(n, NOPE_DIM).contiguous().view(torch.uint8)
rope_bf16_bytes = rope_bf16.contiguous().view(torch.uint8)
scale_bytes = (exponent + 127.0).clamp(0, 255).to(torch.uint8)
arange_nope = torch.arange(NOPE_DIM)
arange_rope = torch.arange(ROPE_DIM * 2)
arange_scale = torch.arange(num_quant_blocks)
token_base = pos_in_block.unsqueeze(-1) * TOKEN_DATA_BYTES
row_idx = block_idx.unsqueeze(-1)
cache[row_idx, token_base + arange_nope] = nope_fp8_bytes
cache[row_idx, token_base + NOPE_DIM + arange_rope] = rope_bf16_bytes
scale_base = (
block_size * TOKEN_DATA_BYTES
+ pos_in_block.unsqueeze(-1) * SCALE_BYTES_PER_TOKEN
)
cache[row_idx, scale_base + arange_scale] = scale_bytes
cache[block_idx, scale_base.squeeze(-1) + num_quant_blocks] = 0
return cache, dequant_ref
def _flash_mla_eager(
q: torch.Tensor,
window_dequant: torch.Tensor,
window_slots: torch.Tensor,
compressed_dequant: torch.Tensor,
compressed_slots: torch.Tensor,
attn_sink: torch.Tensor,
scale: float,
) -> torch.Tensor:
"""MQA attention: gather the (already-dequantized) window+compressed KV
rows referenced by *_slots (-1 = invalid/masked), append a learned
per-head sink logit with zero value contribution, softmax, weighted
sum."""
num_tokens, num_heads, _ = q.shape
q_f32 = q.float()
def gather(dequant, slots):
valid = slots >= 0
safe = slots.clamp(min=0)
kv = dequant[safe] # [T, K, D]
return kv, valid
win_kv, win_valid = gather(window_dequant, window_slots)
if compressed_slots.shape[1] > 0:
comp_kv, comp_valid = gather(compressed_dequant, compressed_slots)
kv = torch.cat([win_kv, comp_kv], dim=1)
valid = torch.cat([win_valid, comp_valid], dim=1)
else:
kv = win_kv
valid = win_valid
scores = torch.einsum("thd,tkd->thk", q_f32, kv) * scale
scores = scores.masked_fill(~valid.unsqueeze(1), float("-inf"))
sink = attn_sink.view(1, num_heads, 1).expand(num_tokens, num_heads, 1)
scores_full = torch.cat([scores, sink], dim=-1)
probs = torch.softmax(scores_full, dim=-1)
out = torch.einsum("thk,tkd->thd", probs[..., :-1], kv)
return out.to(q.dtype)
@pytest.mark.parametrize("num_heads", [1, 64])
@pytest.mark.parametrize("num_tokens", [1, 17])
@pytest.mark.parametrize("swa_only", [True, False])
def test_flash_mla_cpu_kernel_matches_eager(
num_heads: int, num_tokens: int, swa_only: bool
):
torch.manual_seed(0)
dtype = torch.bfloat16
window_block_size = 16
compressed_block_size = 8
num_window = 6
num_compressed = 0 if swa_only else 5
q = torch.randn(num_tokens, num_heads, HEAD_DIM, dtype=dtype)
scale = 1.0 / (HEAD_DIM**0.5)
attn_sink = torch.randn(num_heads, dtype=torch.float32)
num_window_slots = num_tokens * num_window + window_block_size
window_kv_f32 = torch.randn(num_window_slots, HEAD_DIM)
window_cache, window_dequant = _pack_flash_mla_cache(
window_kv_f32, window_block_size
)
window_slots = torch.randint(0, num_window_slots, (num_tokens, num_window))
window_slots[:, 0] = -1 # exercise invalid-slot masking
if num_compressed > 0:
num_compressed_slots = num_tokens * num_compressed + compressed_block_size
compressed_kv_f32 = torch.randn(num_compressed_slots, HEAD_DIM)
compressed_cache, compressed_dequant = _pack_flash_mla_cache(
compressed_kv_f32, compressed_block_size
)
compressed_slots = torch.randint(
0, num_compressed_slots, (num_tokens, num_compressed)
)
if num_compressed > 1:
compressed_slots[:, -1] = -1
else:
compressed_cache = window_cache
compressed_dequant = window_dequant.new_zeros(0, HEAD_DIM)
compressed_slots = window_slots.new_full((num_tokens, 0), -1)
out_ref = _flash_mla_eager(
q,
window_dequant,
window_slots,
compressed_dequant,
compressed_slots,
attn_sink,
scale,
)
out_kernel = q.new_empty(num_tokens, num_heads, HEAD_DIM)
flash_mla_with_kvcache_cpu(
out_kernel,
q,
window_cache,
window_slots,
window_block_size,
compressed_cache,
compressed_slots,
compressed_block_size,
attn_sink,
scale,
)
torch.testing.assert_close(
out_kernel.float(), out_ref.float(), rtol=2e-2, atol=2e-2
)
def test_flash_mla_cpu_kernel_all_invalid_falls_back_to_sink():
"""When every KV slot is masked out, softmax degenerates to the sink
logit alone and the output must be all-zero (no valid value rows)."""
torch.manual_seed(2)
dtype = torch.bfloat16
num_tokens, num_heads, num_window = 3, 4, 5
window_block_size = 16
q = torch.randn(num_tokens, num_heads, HEAD_DIM, dtype=dtype)
scale = 1.0 / (HEAD_DIM**0.5)
attn_sink = torch.randn(num_heads, dtype=torch.float32)
window_kv_f32 = torch.randn(window_block_size, HEAD_DIM)
window_cache, _ = _pack_flash_mla_cache(window_kv_f32, window_block_size)
window_slots = torch.full((num_tokens, num_window), -1, dtype=torch.int64)
compressed_slots = window_slots.new_full((num_tokens, 0), -1)
out = q.new_empty(num_tokens, num_heads, HEAD_DIM)
flash_mla_with_kvcache_cpu(
out,
q,
window_cache,
window_slots,
window_block_size,
window_cache,
compressed_slots,
1,
attn_sink,
scale,
)
assert out.abs().max().item() == 0.0
# ---------------------------------------------------------------------------
# fused_qnorm_rope_kv_insert_cpu (store_cache.cpp)
# ---------------------------------------------------------------------------
def _qnorm_kv_insert_eager(
q, kv, positions, cache, slot_mapping, cos_sin_cache, q_head_padded, eps, bs
):
"""Reference: per-head weight-free RMSNorm + GPT-J RoPE on Q; GPT-J RoPE
+ UE8M0 FP8 quant + paged insert on KV."""
num_quant_blocks = NOPE_DIM // QUANT_BLOCK
num_tokens_full, num_heads_q, _ = q.shape
out_dtype = q.dtype
cos = cos_sin_cache[positions, : ROPE_DIM // 2]
sin = cos_sin_cache[positions, ROPE_DIM // 2 :]
q_f32 = q.float()
var = q_f32.square().mean(dim=-1, keepdim=True)
q_f32 = q_f32 * torch.rsqrt(var + eps)
q_f32[..., NOPE_DIM:] = _apply_gptj_rope(
q_f32[..., NOPE_DIM:], cos.unsqueeze(1), sin.unsqueeze(1)
)
q_out = q.new_zeros(num_tokens_full, q_head_padded, HEAD_DIM)
q_out[:, :num_heads_q, :] = q_f32.to(out_dtype)
num_tokens_insert = slot_mapping.shape[0]
kv_f32 = kv[:num_tokens_insert].float()
kv_f32[..., NOPE_DIM:] = _apply_gptj_rope(
kv_f32[..., NOPE_DIM:], cos[:num_tokens_insert], sin[:num_tokens_insert]
)
kv_bf16 = kv_f32.to(torch.bfloat16)
kv_f32 = kv_bf16.float()
valid = slot_mapping >= 0
if bool(valid.any()):
idx = valid.nonzero(as_tuple=True)[0]
slot_id = slot_mapping[idx]
block_idx = torch.div(slot_id, bs, rounding_mode="floor")
pos_in_block = slot_id % bs
nope = kv_f32[idx, :NOPE_DIM].view(-1, num_quant_blocks, QUANT_BLOCK)
absmax = nope.abs().amax(dim=-1).clamp(min=1e-4)
exponent = torch.ceil(torch.log2(absmax / FP8_MAX))
inv_scale = torch.exp2(-exponent)
scaled = (nope * inv_scale.unsqueeze(-1)).clamp(-FP8_MAX, FP8_MAX)
# Flush fp8-subnormal magnitudes to zero, matching the kernel's
# AVX512 cvt_fp32x16_to_fp8x16, which doesn't preserve exact
# subnormal encodings the way .to(torch.float8_e4m3fn) does.
fp8_min_normal = 2.0**-6
scaled = torch.where(scaled.abs() < fp8_min_normal, 0.0, scaled)
nope_fp8_bytes = (
scaled.to(torch.float8_e4m3fn).view(-1, NOPE_DIM).view(torch.uint8)
)
scale_bytes = (exponent + 127.0).clamp(0, 255).to(torch.uint8)
rope_bf16_bytes = kv_bf16[idx, NOPE_DIM:].contiguous().view(torch.uint8)
arange_nope = torch.arange(NOPE_DIM)
arange_rope = torch.arange(ROPE_DIM * 2)
arange_scale = torch.arange(num_quant_blocks)
row_idx = block_idx.unsqueeze(-1)
token_base = pos_in_block.unsqueeze(-1) * TOKEN_DATA_BYTES
cache[row_idx, token_base + arange_nope] = nope_fp8_bytes
cache[row_idx, token_base + NOPE_DIM + arange_rope] = rope_bf16_bytes
scale_base = (
bs * TOKEN_DATA_BYTES + pos_in_block.unsqueeze(-1) * SCALE_BYTES_PER_TOKEN
)
cache[row_idx, scale_base + arange_scale] = scale_bytes
cache[block_idx, scale_base.squeeze(-1) + num_quant_blocks] = 0
return q_out
@pytest.mark.parametrize("num_tokens", [1, 64])
@pytest.mark.parametrize("n_heads,padded_heads", [(1, 8), (8, 8), (8, 16), (64, 128)])
def test_qnorm_rope_kv_insert_cpu_kernel_matches_eager(
num_tokens: int, n_heads: int, padded_heads: int
):
torch.manual_seed(0)
dtype = torch.bfloat16
eps = 1e-6
max_pos = 4096
block_size = 16
q = torch.randn(num_tokens, n_heads, HEAD_DIM, dtype=dtype)
kv = torch.randn(num_tokens, HEAD_DIM, dtype=dtype)
positions = torch.randint(0, max_pos, (num_tokens,), dtype=torch.int64)
cos_sin_cache = make_cos_sin_cache(max_pos, ROPE_DIM, torch.float32, "cpu")
num_blocks = (num_tokens + block_size - 1) // block_size + 1
slot_mapping = torch.arange(num_tokens, dtype=torch.int64)
if num_tokens > 1:
slot_mapping[-1] = -1 # exercise the "skip invalid slot" branch
cache_ref = _make_cache(num_blocks, block_size)
cache_fused = cache_ref.clone()
q_ref = _qnorm_kv_insert_eager(
q,
kv,
positions,
cache_ref,
slot_mapping,
cos_sin_cache,
padded_heads,
eps,
block_size,
)
q_fused = fused_qnorm_rope_kv_insert_cpu(
q,
kv,
positions,
cache_fused,
slot_mapping,
cos_sin_cache,
padded_heads,
eps,
block_size,
)
torch.testing.assert_close(q_fused.float(), q_ref.float(), rtol=1e-2, atol=1e-2)
if n_heads < padded_heads:
pad_region = q_fused[:, n_heads:padded_heads]
assert pad_region.abs().max().item() == 0.0, (
"padded head slots must be exact zero"
)
_assert_cache_parity(cache_fused, cache_ref, slot_mapping, block_size)
@pytest.mark.parametrize("pad", [1, 5])
def test_qnorm_rope_kv_insert_cpu_kernel_dp_padding(pad: int):
"""slot_mapping shorter than q/kv rows: the KV branch must only touch
the first num_tokens rows while Q-norm+RoPE still runs on all rows."""
torch.manual_seed(1)
dtype = torch.bfloat16
eps = 1e-6
max_pos = 4096
block_size = 16
num_tokens = 17
total = num_tokens + pad
padded_heads = 8
q = torch.randn(total, padded_heads, HEAD_DIM, dtype=dtype)
kv = torch.randn(total, HEAD_DIM, dtype=dtype)
positions = torch.randint(0, max_pos, (total,), dtype=torch.int64)
cos_sin_cache = make_cos_sin_cache(max_pos, ROPE_DIM, torch.float32, "cpu")
num_blocks = (num_tokens + block_size - 1) // block_size + 1
slot_mapping = torch.arange(num_tokens, dtype=torch.int64)
cache_ref = _make_cache(num_blocks, block_size)
cache_fused = cache_ref.clone()
q_ref = _qnorm_kv_insert_eager(
q,
kv,
positions,
cache_ref,
slot_mapping,
cos_sin_cache,
padded_heads,
eps,
block_size,
)
q_fused = fused_qnorm_rope_kv_insert_cpu(
q,
kv,
positions,
cache_fused,
slot_mapping,
cos_sin_cache,
padded_heads,
eps,
block_size,
)
torch.testing.assert_close(q_fused.float(), q_ref.float(), rtol=1e-2, atol=1e-2)
_assert_cache_parity(cache_fused, cache_ref, slot_mapping, block_size)
# ---------------------------------------------------------------------------
# inverse_gptj_rope_o_proj_cpu (store_cache.cpp): DeepseekV4CPUAttention's
# _o_proj de-rotation kernel
# ---------------------------------------------------------------------------
def _inverse_gptj_rope_o_proj_eager(
o: torch.Tensor,
positions: torch.Tensor,
cos_sin_cache: torch.Tensor,
rope_dim: int,
) -> torch.Tensor:
"""NoPE dims pass through as fp32; the RoPE segment is de-rotated by
R(-pos) in GPT-J (interleaved even/odd) convention."""
nope_dim = o.shape[-1] - rope_dim
half_dim = rope_dim // 2
o_f32 = o.float()
out = torch.empty_like(o_f32)
out[..., :nope_dim] = o_f32[..., :nope_dim]
cos = cos_sin_cache[positions, :half_dim].unsqueeze(1)
sin = cos_sin_cache[positions, half_dim:].unsqueeze(1)
e = o_f32[..., nope_dim::2]
od = o_f32[..., nope_dim + 1 :: 2]
out[..., nope_dim::2] = e * cos + od * sin
out[..., nope_dim + 1 :: 2] = od * cos - e * sin
return out
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
@pytest.mark.parametrize("num_tokens", [1, 17])
@pytest.mark.parametrize(
"head_dim,rope_dim,num_heads",
[
(512, 64, 8), # DeepseekV4CPUAttention._o_proj's actual shape
(64, 64, 1), # nope_dim=0: no NoPE region at all
],
)
def test_inverse_gptj_rope_o_proj_cpu_kernel_matches_eager(
dtype: torch.dtype,
num_tokens: int,
head_dim: int,
rope_dim: int,
num_heads: int,
):
torch.manual_seed(0)
max_pos = 4096
o = torch.randn(num_tokens, num_heads, head_dim, dtype=dtype)
positions = torch.randint(0, max_pos, (num_tokens,), dtype=torch.int64)
cos_sin_cache = make_cos_sin_cache(max_pos, rope_dim, torch.float32, "cpu")
out_ref = _inverse_gptj_rope_o_proj_eager(o, positions, cos_sin_cache, rope_dim)
out_fused = inverse_gptj_rope_o_proj_cpu(o, positions, cos_sin_cache, rope_dim)
torch.testing.assert_close(out_fused, out_ref, rtol=1e-2, atol=1e-2)
# ---------------------------------------------------------------------------
# fp8_paged_mqa_logits_cpu (paged_mqa_logits.cpp): sparse indexer's
# decode-path MQA-logits kernel
# ---------------------------------------------------------------------------
def _rand_fp8(*shape: int) -> torch.Tensor:
return torch.randn(shape, dtype=torch.float32).clamp(-3, 3).to(torch.float8_e4m3fn)
def _pages_per_request(block_size: int, seq_lens: list[int]) -> list[int]:
return [-(-seq_len // block_size) for seq_len in seq_lens]
def _write_page(
page_row: torch.Tensor, block_size: int, k_fp8: torch.Tensor, scale: torch.Tensor
) -> None:
"""Write k_fp8 ([block_size, INDEXER_HEAD_DIM]) and scale ([block_size]
fp32) into one page row: K-region then scale-region, matching
compress_norm_rope_store_indexer_cpu's layout."""
k_region = page_row[: block_size * INDEXER_HEAD_DIM]
k_region.copy_(k_fp8.view(torch.uint8).flatten())
scale_region = page_row[block_size * INDEXER_HEAD_DIM :].view(torch.float32)
scale_region[:].copy_(scale)
def _reference_paged_mqa_logits(
q_fp8: torch.Tensor,
weight: torch.Tensor,
ref_k: list[torch.Tensor],
ref_scale: list[torch.Tensor],
max_seq_len: int,
) -> torch.Tensor:
batch_size, _num_heads, _ = q_fp8.shape
out = torch.zeros((batch_size, max_seq_len), dtype=torch.float32)
for b in range(batch_size):
k = ref_k[b].float() # [seq_len, head_dim]
scale = ref_scale[b] # [seq_len]
q = q_fp8[b].float() # [heads, head_dim]
w = weight[b] # [heads]
dot = torch.einsum("jd,hd->jh", k, q).clamp(min=0.0)
score = (dot * w.unsqueeze(0)).sum(dim=1)
out[b, : k.shape[0]] = score * scale
return out
def _build_paged_mqa_logits_batch(
block_size: int,
seq_lens: list[int],
num_heads: int,
pages: torch.Tensor,
seed: int,
):
"""Fill pages (a pre-sized [n_pages, buf_width] uint8 tensor, possibly
non-contiguous in dim 0) with deterministic-per-seed random K/scale
data."""
torch.manual_seed(seed)
batch_size = len(seq_lens)
max_seq_len = max(seq_lens)
per_request = _pages_per_request(block_size, seq_lens)
pages_per_batch = max(per_request)
page_table = torch.full((batch_size, pages_per_batch), -1, dtype=torch.int32)
ref_k: list[torch.Tensor] = []
ref_scale: list[torch.Tensor] = []
next_page = 0
for b, seq_len in enumerate(seq_lens):
k_chunks = []
scale_chunks = []
for p in range(per_request[b]):
page_table[b, p] = next_page
k = _rand_fp8(block_size, INDEXER_HEAD_DIM)
scale = torch.rand(block_size, dtype=torch.float32) * 0.05 + 0.01
_write_page(pages[next_page], block_size, k, scale)
tokens_here = min(block_size, seq_len - p * block_size)
k_chunks.append(k[:tokens_here])
scale_chunks.append(scale[:tokens_here])
next_page += 1
ref_k.append(torch.cat(k_chunks, dim=0))
ref_scale.append(torch.cat(scale_chunks, dim=0))
q_fp8 = _rand_fp8(batch_size, num_heads, INDEXER_HEAD_DIM)
weight = torch.rand(batch_size, num_heads, dtype=torch.float32) * 0.1
seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32)
return q_fp8, weight, seq_lens_t, page_table, ref_k, ref_scale, max_seq_len
def _paged_mqa_logits_page_layout(
block_size: int, seq_lens: list[int]
) -> tuple[int, int]:
n_pages = sum(_pages_per_request(block_size, seq_lens))
buf_width = block_size * INDEXER_HEAD_DIM + block_size * INDEXER_SCALE_BYTES
return n_pages, buf_width
@pytest.mark.parametrize("block_size", [64, 256])
def test_fp8_paged_mqa_logits_cpu_matches_eager(block_size: int):
num_heads = 4
seq_lens = [1, block_size + 7, 3 * block_size]
n_pages, buf_width = _paged_mqa_logits_page_layout(block_size, seq_lens)
pages = torch.zeros((n_pages, buf_width), dtype=torch.uint8)
q_fp8, weight, seq_lens_t, page_table, ref_k, ref_scale, max_seq_len = (
_build_paged_mqa_logits_batch(block_size, seq_lens, num_heads, pages, seed=0)
)
logits = fp8_paged_mqa_logits_cpu(
q_fp8, pages, weight, seq_lens_t, page_table, block_size, max_seq_len
)
ref = _reference_paged_mqa_logits(q_fp8, weight, ref_k, ref_scale, max_seq_len)
for b, seq_len in enumerate(seq_lens):
torch.testing.assert_close(
logits[b, :seq_len], ref[b, :seq_len], atol=1e-2, rtol=1e-2
)
def test_fp8_paged_mqa_logits_cpu_non_contiguous_kv_cache_row_stride():
"""The paged K-cache in production is a per-layer view into a shared
multi-layer allocation, so its page stride generally exceeds the page
byte width -- check the kernel reads a padded/strided buffer identically
to an equivalent tightly-packed one."""
block_size = 64
num_heads = 4
seq_lens = [1, block_size + 7, 3 * block_size]
n_pages, buf_width = _paged_mqa_logits_page_layout(block_size, seq_lens)
cache_pad = 48
tight_pages = torch.zeros((n_pages, buf_width), dtype=torch.uint8)
q_fp8, weight, seq_lens_t, page_table, _, _, max_seq_len = (
_build_paged_mqa_logits_batch(
block_size, seq_lens, num_heads, tight_pages, seed=0
)
)
padded = torch.zeros((n_pages, buf_width + cache_pad), dtype=torch.uint8)
strided_pages = padded[:, :buf_width]
_build_paged_mqa_logits_batch(
block_size, seq_lens, num_heads, strided_pages, seed=0
)
assert strided_pages.stride(0) != strided_pages.size(1)
logits_tight = fp8_paged_mqa_logits_cpu(
q_fp8, tight_pages, weight, seq_lens_t, page_table, block_size, max_seq_len
)
logits_strided = fp8_paged_mqa_logits_cpu(
q_fp8, strided_pages, weight, seq_lens_t, page_table, block_size, max_seq_len
)
# Only [:seq_len] is ever written (and ever read downstream by
# topk_transform_512_cpu) -- positions beyond that are uninitialized, so
# comparing the full row would be comparing unrelated garbage.
for b, seq_len in enumerate(seq_lens):
torch.testing.assert_close(
logits_tight[b, :seq_len], logits_strided[b, :seq_len]
)
# ---------------------------------------------------------------------------
# fused_indexer_q_rope_quant_cpu (indexer.cpp): indexer Q-side RoPE + FP8
# quant, and the dispatcher wiring that routes to it on CPU
# ---------------------------------------------------------------------------
#
# Independent eager reference: this kernel's absmax reduction is over an
# unusual mix (raw NoPE values plus a bf16-round-tripped RoPE half), which
# per_token_group_quant_fp8 does not replicate. Allow <=1 fp8 ULP: the RoPE
# halves round a value sitting exactly on a tie to either side.
INDEXER_Q_N_HEAD = 64
INDEXER_Q_MAX_POS = 4096
def _indexer_q_rope_quant_reference(
positions: torch.Tensor,
q: torch.Tensor,
cos_sin_cache: torch.Tensor,
weights: torch.Tensor,
softmax_scale: float,
head_scale: float,
) -> tuple[torch.Tensor, torch.Tensor]:
_num_tokens, _num_heads, head_dim = q.shape
half_rot_dim = cos_sin_cache.shape[-1] // 2
rot_dim = 2 * half_rot_dim
nope_dim = head_dim - rot_dim
q_f32 = q.float()
cos_sin_f32 = cos_sin_cache.float()
pos_cos_sin = cos_sin_f32[positions] # [T, rot_dim]
cos = pos_cos_sin[:, :half_rot_dim].unsqueeze(1)
sin = pos_cos_sin[:, half_rot_dim:].unsqueeze(1)
nope = q_f32[..., :nope_dim]
rot = q_f32[..., nope_dim:]
x_even = rot[..., 0::2]
x_odd = rot[..., 1::2]
r_even = (x_even * cos - x_odd * sin).to(torch.bfloat16).to(torch.float32)
r_odd = (x_odd * cos + x_even * sin).to(torch.bfloat16).to(torch.float32)
rope_out = torch.stack((r_even, r_odd), dim=-1).flatten(-2)
combined = torch.cat((nope, rope_out), dim=-1) # [T, H, head_dim]
amax = combined.abs().amax(dim=-1).clamp(min=1e-4) # [T, H]
exponent = torch.ceil(torch.log2(amax / FP8_MAX))
scale = torch.exp2(exponent)
inv_scale = scale.reciprocal().unsqueeze(-1)
q_fp8 = (combined * inv_scale).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn)
weights_out = weights.to(torch.float32) * scale * softmax_scale * head_scale
return q_fp8, weights_out
@pytest.mark.parametrize("num_tokens", [1, 257])
def test_fused_indexer_q_rope_quant_cpu_matches_eager(num_tokens: int):
torch.manual_seed(0)
q = torch.randn(num_tokens, INDEXER_Q_N_HEAD, INDEXER_HEAD_DIM, dtype=torch.float32)
positions = torch.randint(0, INDEXER_Q_MAX_POS, (num_tokens,), dtype=torch.int64)
cos_sin_cache = torch.randn(INDEXER_Q_MAX_POS, ROPE_DIM, dtype=torch.float32)
weights = torch.randn(num_tokens, INDEXER_Q_N_HEAD, dtype=torch.float32)
softmax_scale = INDEXER_HEAD_DIM**-0.5
head_scale = INDEXER_Q_N_HEAD**-0.5
q_fp8_ref, weights_ref = _indexer_q_rope_quant_reference(
positions, q, cos_sin_cache, weights, softmax_scale, head_scale
)
q_fp8_fused = torch.empty_like(q, dtype=torch.float8_e4m3fn)
weights_fused = torch.empty_like(weights, dtype=torch.float32)
fused_indexer_q_rope_quant_cpu(
positions,
q,
cos_sin_cache,
q_fp8_fused,
weights,
softmax_scale,
head_scale,
weights_fused,
)
max_ulp = int(fp8_ulp_distance(q_fp8_fused, q_fp8_ref).max().item())
assert max_ulp <= 1, f"q_fp8 differs by {max_ulp} ULP (>1)"
torch.testing.assert_close(weights_fused, weights_ref, rtol=1e-5, atol=1e-6)
def test_fused_indexer_q_rope_quant_dispatches_to_cpu_kernel():
"""End-to-end wiring check: the production dispatcher must route to the
CPU op on a CPU tensor, not silently fall through to the Triton kernel
(which would also "work" under triton-cpu, masking a dispatch
regression). The dispatcher does its CPU import inside the CPU branch,
so the patch target is the source module."""
torch.manual_seed(0)
num_tokens = 1
q = torch.randn(num_tokens, INDEXER_Q_N_HEAD, INDEXER_HEAD_DIM, dtype=torch.float32)
positions = torch.randint(0, INDEXER_Q_MAX_POS, (num_tokens,), dtype=torch.int64)
cos_sin_cache = torch.randn(INDEXER_Q_MAX_POS, ROPE_DIM, dtype=torch.float32)
weights = torch.randn(num_tokens, INDEXER_Q_N_HEAD, dtype=torch.float32)
softmax_scale = INDEXER_HEAD_DIM**-0.5
head_scale = INDEXER_Q_N_HEAD**-0.5
q_fp8_ref, weights_ref = _indexer_q_rope_quant_reference(
positions, q, cos_sin_cache, weights, softmax_scale, head_scale
)
with mock.patch("vllm._custom_ops.fused_indexer_q_rope_quant_cpu") as mocked:
fused_indexer_q_rope_quant(
positions,
q.clone(),
cos_sin_cache,
weights,
softmax_scale,
head_scale,
use_fp4=False,
)
assert mocked.called, (
"fused_indexer_q_rope_quant did not dispatch to the CPU kernel "
"on a CPU tensor"
)
q_fp8_fused, weights_fused = fused_indexer_q_rope_quant(
positions,
q.clone(),
cos_sin_cache,
weights,
softmax_scale,
head_scale,
use_fp4=False,
)
max_ulp = int(fp8_ulp_distance(q_fp8_fused, q_fp8_ref).max().item())
assert max_ulp <= 1, f"q_fp8 differs by {max_ulp} ULP (>1)"
torch.testing.assert_close(weights_fused, weights_ref, rtol=1e-5, atol=1e-6)
# ---------------------------------------------------------------------------
# mhc_pre_cpu / mhc_post_cpu / hc_head_fused_cpu (mhc.cpp): CPU ports of the
# mHC gating kernels
# ---------------------------------------------------------------------------
def _rmsnorm_nw_torch(x: torch.Tensor, eps: float) -> torch.Tensor:
"""Weight-free RMSNorm over the last dim, matching hc_head_fused_cpu."""
x_f32 = x.float()
var = x_f32.square().mean(dim=-1, keepdim=True)
return (x_f32 * torch.rsqrt(var + eps)).to(x.dtype)
def _hc_head_torch(
hidden_states: torch.Tensor,
hc_fn: torch.Tensor,
hc_scale: torch.Tensor,
hc_base: torch.Tensor,
rms_norm_eps: float,
hc_eps: float,
) -> torch.Tensor:
"""Eager reference for hc_head_fused_cpu: round-trips the RMSNorm
output through bf16 before the linear, matching the kernel's bf16
arithmetic."""
hc_mult, hidden_size = hidden_states.shape[-2:]
outer_shape = hidden_states.shape[:-2]
hs_flat = hidden_states.reshape(-1, hc_mult, hidden_size)
x_flat = hs_flat.flatten(-2)
x_normed = _rmsnorm_nw_torch(x_flat, rms_norm_eps)
mixes = torch.nn.functional.linear(x_normed.float(), hc_fn)
pre = torch.sigmoid(mixes * hc_scale + hc_base) + hc_eps
out = torch.sum(pre.unsqueeze(-1) * hs_flat.float(), dim=1).to(hidden_states.dtype)
return out.view(*outer_shape, hidden_size)
@pytest.mark.parametrize("num_tokens", [1, 128])
@pytest.mark.parametrize("hidden_size", [4096, 7168])
def test_mhc_pre_cpu(num_tokens: int, hidden_size: int):
"""mhc_pre_cpu vs. the mhc_pre_torch eager reference it replaces."""
torch.set_default_device(DEVICE)
set_random_seed(0)
hc_mult = 4
hc_mult3 = 2 * hc_mult + hc_mult * hc_mult
residual = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16)
fn = torch.randn((hc_mult3, hc_mult * hidden_size), dtype=torch.float32) * 1e-4
hc_scale = torch.randn((3,), dtype=torch.float32) * 0.1
hc_base = torch.randn((hc_mult3,), dtype=torch.float32) * 0.1
rms_eps = hc_eps = 1e-6
hc_post_mult_value = 2.0
sinkhorn_repeat = 8
ref = mhc_kernels.mhc_pre_torch(
residual,
fn,
hc_scale,
hc_base,
rms_eps,
hc_eps,
hc_eps,
hc_post_mult_value,
sinkhorn_repeat,
)
out = mhc_kernels.mhc_pre_cpu(
residual,
fn,
hc_scale,
hc_base,
rms_eps,
hc_eps,
hc_eps,
hc_post_mult_value,
sinkhorn_repeat,
)
for actual, expected in zip(out, ref, strict=True):
torch.testing.assert_close(actual, expected, atol=5e-2, rtol=1e-2)
@pytest.mark.parametrize("num_tokens", [1, 128])
@pytest.mark.parametrize("hidden_size", [4096, 7168])
def test_mhc_post_cpu(num_tokens: int, hidden_size: int):
"""mhc_post_cpu vs. the mhc_post_torch eager reference it replaces."""
torch.set_default_device(DEVICE)
set_random_seed(0)
hc_mult = 4
x = torch.randn((num_tokens, hidden_size), dtype=torch.bfloat16)
residual = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16)
post_layer_mix = torch.randn((num_tokens, hc_mult, 1), dtype=torch.float32)
comb_res_mix = torch.randn((num_tokens, hc_mult, hc_mult), dtype=torch.float32)
ref = mhc_kernels.mhc_post_torch(x, residual, post_layer_mix, comb_res_mix)
out = mhc_kernels.mhc_post_cpu(x, residual, post_layer_mix, comb_res_mix)
torch.testing.assert_close(out, ref, atol=5e-2, rtol=1e-2)
@pytest.mark.parametrize("num_tokens", [1, 128])
@pytest.mark.parametrize("hidden_size", [4096, 7168])
def test_hc_head_cpu(num_tokens: int, hidden_size: int):
"""hc_head_fused_cpu vs. the _hc_head_torch eager reference it
replaces."""
torch.set_default_device(DEVICE)
set_random_seed(0)
hc_mult = 4
residual = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16)
fn = torch.randn((hc_mult, hc_mult * hidden_size), dtype=torch.float32) * 1e-4
hc_scale = torch.randn((1,), dtype=torch.float32) * 0.1
hc_base = torch.randn((hc_mult,), dtype=torch.float32) * 0.1
rms_eps = hc_eps = 1e-6
ref = _hc_head_torch(residual, fn, hc_scale, hc_base, rms_eps, hc_eps)
out = mhc_kernels.hc_head_fused_cpu(
residual, fn, hc_scale, hc_base, rms_eps, hc_eps
)
torch.testing.assert_close(out, ref, atol=5e-2, rtol=1e-2)
Reference in New Issue View Git Blame Copy Permalink