| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Unit tests for the fused inverse RoPE + block-scaled FP8 quantization kernel. |
| 4 | |
| 5 | Tests compare the fused kernel against a reference implementation built from |
| 6 | the existing separate operations (inverse RoPE via rotate_neox + FP8 quant |
| 7 | via per_token_group_quant_fp8). |
| 8 | |
| 9 | The reference faithfully reproduces the exact flow in |
| 10 | deepseek_v4/attention.py:295-310: |
| 11 | 1. Apply inverse RoPE (NeoX style, last rope_dim=64 dims of each head) |
| 12 | 2. Reshape [T, H, head_dim] -> [T, G, D] |
| 13 | 3. Transpose+flatten to [G*T, D], quantize, reshape back |
| 14 | 4. Return o_fp8 and o_scale with strides (D, T*D, 1) and (S, T*S, 1) |
| 15 | (non-contiguous [T, G, ...] view backed by contiguous [G, T, ...] memory) |
| 16 | |
| 17 | Usage: |
| 18 | pytest tests/kernels/test_fused_inv_rope_fp8_quant.py -v |
| 19 | """ |
| 20 | |
| 21 | import pytest |
| 22 | import torch |
| 23 | |
| 24 | from vllm.model_executor.layers.quantization.utils.fp8_utils import ( |
| 25 | deepgemm_post_process_fp8_weight_block, |
| 26 | ) |
| 27 | from vllm.models.deepseek_v4.common.ops import fused_inv_rope_fp8_quant |
| 28 | from vllm.platforms import current_platform |
| 29 | from vllm.utils.deep_gemm import fp8_einsum |
| 30 | |
| 31 | # -- Default dimensions matching DeepSeek V3/V4 -------------------------- |
| 32 | HEAD_DIM = 512 |
| 33 | NOPE_DIM = 448 |
| 34 | ROPE_DIM = 64 |
| 35 | QUANT_GROUP_SIZE = 128 |
| 36 | FP8_MAX = 448.0 # torch.finfo(torch.float8_e4m3fn).max |
| 37 | FP8_DTYPE = torch.float8_e4m3fn |
| 38 | EPS = 1e-10 |
| 39 | |
| 40 | |
| 41 | @pytest.mark.skipif( |
| 42 | not current_platform.is_device_capability_family(100), reason="Requires SM100/SM103" |
| 43 | ) |
| 44 | @pytest.mark.parametrize("groups", [2, 4]) |
| 45 | @pytest.mark.parametrize("tokens", [1, 6, 18, 32]) |
| 46 | def test_dsv41_fused_wo_a_matches_deepgemm(groups, tokens): |
| 47 | """Fused WO-A matches inverse RoPE + fp8_einsum + MXFP8 quantize.""" |
| 48 | from flashinfer import mxfp8_quantize |
| 49 | |
| 50 | from vllm.models.deepseek_v41.nvidia.ops.fused_wo_a import _FUSED_WO_A_KERNEL |
| 51 | |
| 52 | torch.manual_seed(0) |
| 53 | n = groups * 1024 |
| 54 | x = torch.randn(tokens, 64, 512, device="cuda", dtype=torch.bfloat16)[:, : n // 128] |
| 55 | positions = torch.randint(0, 128, (tokens,), device="cuda") |
| 56 | rope = make_cos_sin_cache(128, device="cuda") |
| 57 | wq = torch.randn(n, 4096, device="cuda").to(FP8_DTYPE) |
| 58 | scale = torch.exp2(torch.randint(-8, 0, (n, 128), device="cuda").float()) |
| 59 | wq, ws = deepgemm_post_process_fp8_weight_block( |
| 60 | wq, scale, (1, 32), False, True, groups |
| 61 | ) |
| 62 | q, sf = _FUSED_WO_A_KERNEL( |
| 63 | x=x, positions=positions, rope=rope, weight=wq, weight_scale=ws |
| 64 | ) |
| 65 | |
| 66 | a, sa = fused_inv_rope_fp8_quant( |
| 67 | x, positions, rope, groups, 8, quant_group_size=32, tma_aligned_scales=True |
| 68 | ) |
| 69 | y = torch.empty(tokens, groups, 1024, device="cuda", dtype=torch.bfloat16) |
| 70 | fp8_einsum("bhr,hdr->bhd", (a, sa), (wq, ws), y, recipe=(1, 1, 32)) |
| 71 | ref, ref_sf = mxfp8_quantize(y.flatten(1), backend="cute-dsl") |
| 72 | |
| 73 | def rows(scales): # F8_128x4 bytes -> [128 rows, n / 32 blocks] |
| 74 | return scales.view(n // 128, 32, 4, 4).permute(2, 1, 0, 3).reshape(128, -1) |
| 75 | |
| 76 | actual, expected = ( |
| 77 | d.float().view(tokens, -1, 32) * torch.exp2(rows(s)[:tokens, :, None] - 127.0) |
| 78 | for d, s in ((q, sf), (ref, ref_sf)) |
| 79 | ) |
| 80 | assert (actual - expected).norm() / expected.norm() < 3e-3 |
| 81 | assert not rows(sf)[tokens:].any(), "padded scale rows must be zero" |
| 82 | |
| 83 | |
| 84 | # ========================================================================= |
| 85 | # Helpers |
| 86 | # ========================================================================= |
| 87 | |
| 88 | |
| 89 | def assert_dequant_close( |
| 90 | fp8_a: torch.Tensor, |
| 91 | scale_a: torch.Tensor, |
| 92 | fp8_b: torch.Tensor, |
| 93 | scale_b: torch.Tensor, |
| 94 | msg: str = "", |
| 95 | ): |
| 96 | """Compare two FP8-quantized tensors via their dequantized values. |
| 97 | |
| 98 | Uses cosine-similarity-based diff (same as deep_gemm calc_diff). |
| 99 | Both fused and reference paths rotate in fp32 using an fp32 |
| 100 | cos_sin_cache, so differences are only fp32 ordering ULPs that can |
| 101 | occasionally shift FP8 values at quantization boundaries. |
| 102 | """ |
| 103 | S = scale_a.shape[-1] |
| 104 | shape = fp8_a.shape |
| 105 | |
| 106 | dq_a = fp8_a.float() * scale_a.unsqueeze(-1).expand( |
| 107 | *shape[:-1], S, QUANT_GROUP_SIZE |
| 108 | ).reshape(shape) |
| 109 | dq_b = fp8_b.float() * scale_b.unsqueeze(-1).expand( |
| 110 | *shape[:-1], S, QUANT_GROUP_SIZE |
| 111 | ).reshape(shape) |
| 112 | |
| 113 | # Cosine diff: 1 - cos_sim (0 = identical, higher = worse) |
| 114 | dq_a_flat = dq_a.flatten().float() |
| 115 | dq_b_flat = dq_b.flatten().float() |
| 116 | cos_sim = torch.nn.functional.cosine_similarity( |
| 117 | dq_a_flat.unsqueeze(0), dq_b_flat.unsqueeze(0) |
| 118 | ).item() |
| 119 | diff = 1.0 - cos_sim |
| 120 | |
| 121 | assert diff < 1e-4, f"Dequant diff too large: {diff:.8f} (expected < 1e-4). {msg}" |
| 122 | |
| 123 | |
| 124 | def rotate_gptj(x: torch.Tensor) -> torch.Tensor: |
| 125 | """GPT-J style rotation: interleaved pairs, negate-swap. |
| 126 | |
| 127 | Matches vllm/model_executor/layers/rotary_embedding/common.py:23-27. |
| 128 | DeepseekV4 uses is_neox_style=False, so this is the correct rotation. |
| 129 | """ |
| 130 | x1 = x[..., ::2] |
| 131 | x2 = x[..., 1::2] |
| 132 | x = torch.stack((-x2, x1), dim=-1) |
| 133 | return x.flatten(-2) |
| 134 | |
| 135 | |
| 136 | def make_cos_sin_cache( |
| 137 | max_pos: int, |
| 138 | rope_dim: int = ROPE_DIM, |
| 139 | dtype: torch.dtype = torch.float32, |
| 140 | device: str = "cuda", |
| 141 | ) -> torch.Tensor: |
| 142 | """Create a synthetic cos_sin_cache matching the layout used by |
| 143 | DeepseekV4ScalingRotaryEmbedding._compute_cos_sin_cache. |
| 144 | |
| 145 | Shape: [max_pos, rope_dim] where first half is cos, second half is sin. |
| 146 | The fused kernel requires fp32; callers can override dtype if passing |
| 147 | the cache into the bf16-only paths. |
| 148 | """ |
| 149 | half = rope_dim // 2 |
| 150 | # Use random but bounded frequencies so cos/sin are well-behaved |
| 151 | inv_freq = 1.0 / ( |
| 152 | 10000.0 ** (torch.arange(0, half, device=device, dtype=torch.float32) / half) |
| 153 | ) |
| 154 | t = torch.arange(max_pos, device=device, dtype=torch.float32) |
| 155 | freqs = torch.outer(t, inv_freq) # [max_pos, half] |
| 156 | cos = freqs.cos() |
| 157 | sin = freqs.sin() |
| 158 | cache = torch.cat((cos, sin), dim=-1) # [max_pos, rope_dim] |
| 159 | return cache.to(dtype) |
| 160 | |
| 161 | |
| 162 | def reference_inv_rope( |
| 163 | o: torch.Tensor, |
| 164 | positions: torch.Tensor, |
| 165 | cos_sin_cache: torch.Tensor, |
| 166 | nope_dim: int = NOPE_DIM, |
| 167 | rope_dim: int = ROPE_DIM, |
| 168 | ) -> torch.Tensor: |
| 169 | """Apply inverse RoPE to the last rope_dim dimensions of each head. |
| 170 | |
| 171 | Matches the GPT-J inverse rotation in pos_encoding_kernels.cu, which |
| 172 | promotes the cache to fp32 and performs the rotation in fp32. The |
| 173 | result is cast back to the input dtype. |
| 174 | |
| 175 | Args: |
| 176 | o: [T, H, head_dim] bf16 |
| 177 | positions: [T] int64 |
| 178 | cos_sin_cache: [max_pos, rope_dim] fp32 |
| 179 | |
| 180 | Returns: |
| 181 | o with inverse RoPE applied on the rope portion (bf16). |
| 182 | |
| 183 | """ |
| 184 | assert cos_sin_cache.dtype == torch.float32 |
| 185 | cos_sin = cos_sin_cache[positions] # [T, rope_dim] fp32 |
| 186 | half = rope_dim // 2 |
| 187 | cos = cos_sin[:, :half] |
| 188 | sin = cos_sin[:, half:] |
| 189 | |
| 190 | # GPT-J style: repeat_interleave (not repeat) to match interleaved pairs |
| 191 | cos = cos.repeat_interleave(2, dim=-1).unsqueeze(1) |
| 192 | sin = sin.repeat_interleave(2, dim=-1).unsqueeze(1) |
| 193 | sin = -sin # inverse |
| 194 | |
| 195 | o_pass = o[..., :nope_dim] |
| 196 | o_rot_f32 = o[..., nope_dim:].float() |
| 197 | o_rot_f32 = o_rot_f32 * cos + rotate_gptj(o_rot_f32) * sin |
| 198 | o_rot = o_rot_f32.to(o.dtype) |
| 199 | |
| 200 | return torch.cat([o_pass, o_rot], dim=-1) |
| 201 | |
| 202 | |
| 203 | def _ref_ue8m0_quant_block(x_f32: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| 204 | """Per-block UE8M0 FP8 quantization in pure float32. |
| 205 | |
| 206 | Matches the Triton kernel logic exactly: |
| 207 | absmax -> 2^ceil(log2(absmax / fp8_max)) -> clamp(x / scale) -> fp8 |
| 208 | |
| 209 | Args: |
| 210 | x_f32: [..., quant_group_size] float32 quantization blocks. |
| 211 | |
| 212 | Returns: |
| 213 | x_fp8: same shape, float8_e4m3fn |
| 214 | scales: [...] float32, one scale per block |
| 215 | |
| 216 | """ |
| 217 | absmax = x_f32.abs().amax(dim=-1, keepdim=True).clamp(min=EPS) |
| 218 | scale_raw = absmax * (1.0 / FP8_MAX) |
| 219 | scale = torch.exp2(torch.ceil(torch.log2(scale_raw))) |
| 220 | x_scaled = (x_f32 / scale).clamp(-FP8_MAX, FP8_MAX) |
| 221 | x_fp8 = x_scaled.to(FP8_DTYPE) |
| 222 | return x_fp8, scale.squeeze(-1) |
| 223 | |
| 224 | |
| 225 | def reference_inv_rope_fp8_quant( |
| 226 | o: torch.Tensor, |
| 227 | positions: torch.Tensor, |
| 228 | cos_sin_cache: torch.Tensor, |
| 229 | n_groups: int, |
| 230 | heads_per_group: int, |
| 231 | nope_dim: int = NOPE_DIM, |
| 232 | rope_dim: int = ROPE_DIM, |
| 233 | quant_group_size: int = QUANT_GROUP_SIZE, |
| 234 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 235 | """Full reference: inverse RoPE in fp32 + UE8M0 FP8 quant in fp32. |
| 236 | |
| 237 | Mimics the Triton kernel's precision path exactly: |
| 238 | Load bf16 -> cast to fp32 -> apply inverse RoPE with fp32 cos/sin -> |
| 239 | UE8M0 quant in fp32 -> write fp8 + scale |
| 240 | |
| 241 | Returns: |
| 242 | o_fp8: [T, G, D] FP8 with strides (D, T*D, 1) |
| 243 | o_scale: [T, G, S] FP32 with strides (S, T*S, 1) |
| 244 | |
| 245 | """ |
| 246 | assert cos_sin_cache.dtype == torch.float32 |
| 247 | T, _H, head_dim = o.shape |
| 248 | d = heads_per_group * head_dim |
| 249 | S = d // quant_group_size |
| 250 | half_rope = rope_dim // 2 |
| 251 | chunks_per_head = head_dim // quant_group_size |
| 252 | |
| 253 | # Reshape [T, H, head_dim] -> [T, G, heads_per_group, head_dim] |
| 254 | o_4d = o.view(T, n_groups, heads_per_group, head_dim) |
| 255 | |
| 256 | # Lookup cos/sin directly in fp32 |
| 257 | cos_sin = cos_sin_cache[positions] # [T, rope_dim] fp32 |
| 258 | cos = cos_sin[:, :half_rope] # [T, half_rope] fp32 |
| 259 | sin = cos_sin[:, half_rope:] # [T, half_rope] fp32 |
| 260 | |
| 261 | # Allocate outputs in [G, T, ...] contiguous layout |
| 262 | fp8_buf = torch.empty(n_groups, T, d, dtype=FP8_DTYPE, device=o.device) |
| 263 | scale_buf = torch.empty(n_groups, T, S, dtype=torch.float32, device=o.device) |
| 264 | |
| 265 | # Process each quant block, matching the Triton kernel's per-program logic |
| 266 | for g in range(n_groups): |
| 267 | for qb in range(S): |
| 268 | head_in_group = qb // chunks_per_head |
| 269 | chunk_in_head = qb % chunks_per_head |
| 270 | offset = chunk_in_head * quant_group_size |
| 271 | |
| 272 | # Load 128 bf16 elements and promote to fp32 for rotation+quant |
| 273 | block = o_4d[:, g, head_in_group, offset : offset + quant_group_size] |
| 274 | x = block.float() |
| 275 | |
| 276 | # Apply inverse RoPE in fp32 if this is the last chunk |
| 277 | # GPT-J style: interleaved pairs (even=x, odd=y) |
| 278 | if chunk_in_head == chunks_per_head - 1: |
| 279 | rope_start = nope_dim % quant_group_size # 64 |
| 280 | rope_region = x[:, rope_start:].clone() |
| 281 | x_vals = rope_region[:, ::2] |
| 282 | y_vals = rope_region[:, 1::2] |
| 283 | x_new = x_vals * cos + y_vals * sin |
| 284 | y_new = y_vals * cos - x_vals * sin |
| 285 | x = x.clone() |
| 286 | x[:, rope_start::2] = x_new |
| 287 | x[:, rope_start + 1 :: 2] = y_new |
| 288 | |
| 289 | # UE8M0 quant in fp32 |
| 290 | x_fp8, scale = _ref_ue8m0_quant_block(x) |
| 291 | |
| 292 | # Write to [G, T, D] contiguous memory |
| 293 | fp8_buf[g, :, qb * quant_group_size : (qb + 1) * quant_group_size] = x_fp8 |
| 294 | scale_buf[g, :, qb] = scale |
| 295 | |
| 296 | # Return transposed views |
| 297 | return fp8_buf.transpose(0, 1), scale_buf.transpose(0, 1) |
| 298 | |
| 299 | |
| 300 | # ========================================================================= |
| 301 | # Tests |
| 302 | # ========================================================================= |
| 303 | |
| 304 | |
| 305 | @pytest.mark.parametrize("num_tokens", [1, 7, 128, 4096]) |
| 306 | @pytest.mark.parametrize("quant_group_size", [32, 128]) |
| 307 | @torch.inference_mode() |
| 308 | def test_packed_quantization_preserves_rotary_dimensions(num_tokens, quant_group_size): |
| 309 | """RoPE may span multiple scale blocks; pack every scale and padding row.""" |
| 310 | from vllm.models.deepseek_v4.common.ops.fused_inv_rope_fp8_quant import ( |
| 311 | fused_inv_rope_fp8_quant as quantize, |
| 312 | ) |
| 313 | |
| 314 | torch.manual_seed(42) |
| 315 | groups, heads_per_group = 2, 2 |
| 316 | x = torch.randn( |
| 317 | num_tokens, |
| 318 | groups * heads_per_group, |
| 319 | HEAD_DIM, |
| 320 | device="cuda", |
| 321 | dtype=torch.bfloat16, |
| 322 | ) |
| 323 | cache = make_cos_sin_cache(4096) |
| 324 | positions = torch.randint(4096, (num_tokens,), device="cuda") |
| 325 | q, packed_scales = quantize( |
| 326 | x, |
| 327 | positions, |
| 328 | cache, |
| 329 | groups, |
| 330 | heads_per_group, |
| 331 | quant_group_size=quant_group_size, |
| 332 | tma_aligned_scales=True, |
| 333 | ) |
| 334 | scale_bytes = torch.stack( |
| 335 | [(packed_scales >> (8 * i)) & 0xFF for i in range(4)], dim=-1 |
| 336 | ).flatten(-2) |
| 337 | scales = torch.exp2(scale_bytes.float() - 127) |
| 338 | rotated = reference_inv_rope(x.float(), positions, cache) |
| 339 | ref_q, ref_scales = _ref_ue8m0_quant_block( |
| 340 | rotated.reshape(num_tokens, groups, -1, quant_group_size) |
| 341 | ) |
| 342 | torch.testing.assert_close(scales, ref_scales, rtol=0, atol=0) |
| 343 | actual = q.float().unflatten(-1, (-1, quant_group_size)) * scales.unsqueeze(-1) |
| 344 | expected = ref_q.float() * ref_scales.unsqueeze(-1) |
| 345 | assert (actual - expected).norm() / expected.norm() < 1e-3 |
| 346 | |
| 347 | aligned_tokens = (num_tokens + 3) // 4 * 4 |
| 348 | full_scales = packed_scales.transpose(0, 1).as_strided( |
| 349 | (groups, aligned_tokens, packed_scales.shape[-1]), |
| 350 | (packed_scales.stride(1), 1, packed_scales.stride(2)), |
| 351 | ) |
| 352 | assert torch.count_nonzero(full_scales[:, num_tokens:]) == 0 |
| 353 | |
| 354 | |
| 355 | @pytest.mark.parametrize("num_tokens", [1, 7, 32, 128]) |
| 356 | @pytest.mark.parametrize( |
| 357 | "num_heads,n_groups", |
| 358 | [(64, 8), (32, 4), (128, 8)], |
| 359 | ids=["H64_G8", "H32_G4", "H128_G8"], |
| 360 | ) |
| 361 | @pytest.mark.parametrize("seed", [0, 42]) |
| 362 | @pytest.mark.parametrize("quantize", [True, False], ids=["fp8", "bf16"]) |
| 363 | @torch.inference_mode() |
| 364 | def test_correctness(num_tokens, num_heads, n_groups, seed, quantize): |
| 365 | """Compare fused kernel against reference for FP8 values and scales.""" |
| 366 | torch.manual_seed(seed) |
| 367 | |
| 368 | heads_per_group = num_heads // n_groups |
| 369 | max_pos = 4096 |
| 370 | device = "cuda" |
| 371 | |
| 372 | # Create inputs |
| 373 | o = torch.randn( |
| 374 | num_tokens, num_heads, HEAD_DIM, device=device, dtype=torch.bfloat16 |
| 375 | ) |
| 376 | positions = torch.randint( |
| 377 | 0, max_pos, (num_tokens,), device=device, dtype=torch.long |
| 378 | ) |
| 379 | cos_sin_cache = make_cos_sin_cache( |
| 380 | max_pos, ROPE_DIM, dtype=torch.float32, device=device |
| 381 | ) |
| 382 | |
| 383 | if not quantize: |
| 384 | ref_output = reference_inv_rope(o, positions, cos_sin_cache).view( |
| 385 | num_tokens, n_groups, -1 |
| 386 | ) |
| 387 | fused_output, fused_scale = fused_inv_rope_fp8_quant( |
| 388 | o.clone(), |
| 389 | positions, |
| 390 | cos_sin_cache, |
| 391 | n_groups, |
| 392 | heads_per_group, |
| 393 | quantize=False, |
| 394 | ) |
| 395 | d = heads_per_group * HEAD_DIM |
| 396 | assert fused_output.shape == (num_tokens, n_groups, d) |
| 397 | assert fused_scale.numel() == 0 |
| 398 | assert fused_output.dtype == torch.bfloat16 |
| 399 | assert fused_output.stride() == (d, num_tokens * d, 1) |
| 400 | torch.testing.assert_close(fused_output, ref_output, rtol=1e-2, atol=1e-2) |
| 401 | return |
| 402 | |
| 403 | # Reference |
| 404 | ref_fp8, ref_scale = reference_inv_rope_fp8_quant( |
| 405 | o.clone(), |
| 406 | positions, |
| 407 | cos_sin_cache, |
| 408 | n_groups, |
| 409 | heads_per_group, |
| 410 | ) |
| 411 | |
| 412 | # Fused kernel |
| 413 | fused_fp8, fused_scale = fused_inv_rope_fp8_quant( |
| 414 | o.clone(), |
| 415 | positions, |
| 416 | cos_sin_cache, |
| 417 | n_groups, |
| 418 | heads_per_group, |
| 419 | ) |
| 420 | |
| 421 | # Check shapes |
| 422 | d = heads_per_group * HEAD_DIM |
| 423 | S = d // QUANT_GROUP_SIZE |
| 424 | assert ref_fp8.shape == (num_tokens, n_groups, d) |
| 425 | assert fused_fp8.shape == (num_tokens, n_groups, d) |
| 426 | assert ref_scale.shape == (num_tokens, n_groups, S) |
| 427 | assert fused_scale.shape == (num_tokens, n_groups, S) |
| 428 | |
| 429 | # Scales: exact match (both use identical UE8M0 algorithm) |
| 430 | # Scales may differ by one UE8M0 step (factor of 2) if fp32 rotation |
| 431 | # ordering shifts absmax across a power-of-2 boundary. Check ratio is |
| 432 | # close to 1. |
| 433 | scale_ratio = fused_scale / ref_scale.clamp(min=1e-30) |
| 434 | assert scale_ratio.max() <= 2.0 and scale_ratio.min() >= 0.5, ( |
| 435 | f"Scale ratio out of [0.5, 2]: min={scale_ratio.min():.4f} " |
| 436 | f"max={scale_ratio.max():.4f}" |
| 437 | ) |
| 438 | |
| 439 | # Compare via dequant (Triton vs PyTorch fp32 may differ by ULPs) |
| 440 | assert_dequant_close(ref_fp8, ref_scale, fused_fp8, fused_scale) |
| 441 | |
| 442 | |
| 443 | @pytest.mark.parametrize("num_tokens", [1, 7, 32, 128]) |
| 444 | @pytest.mark.parametrize( |
| 445 | "num_heads,n_groups", |
| 446 | [(64, 8), (128, 8)], |
| 447 | ids=["H64_G8", "H128_G8"], |
| 448 | ) |
| 449 | @torch.inference_mode() |
| 450 | def test_output_strides(num_tokens, num_heads, n_groups): |
| 451 | """Verify fused output layout: |
| 452 | - FP8: logical [T, G, D] backed by contiguous [G, T, D]. |
| 453 | - Scale: MN-major TMA-aligned (column-major: T-stride=1). |
| 454 | """ |
| 455 | heads_per_group = num_heads // n_groups |
| 456 | max_pos = 4096 |
| 457 | device = "cuda" |
| 458 | |
| 459 | o = torch.randn( |
| 460 | num_tokens, num_heads, HEAD_DIM, device=device, dtype=torch.bfloat16 |
| 461 | ) |
| 462 | positions = torch.randint( |
| 463 | 0, max_pos, (num_tokens,), device=device, dtype=torch.long |
| 464 | ) |
| 465 | cos_sin_cache = make_cos_sin_cache(max_pos, device=device) |
| 466 | |
| 467 | fused_fp8, fused_scale = fused_inv_rope_fp8_quant( |
| 468 | o.clone(), |
| 469 | positions, |
| 470 | cos_sin_cache, |
| 471 | n_groups, |
| 472 | heads_per_group, |
| 473 | ) |
| 474 | |
| 475 | # FP8: logical [T, G, D] backed by [G, T, D] row-major |
| 476 | d = heads_per_group * HEAD_DIM |
| 477 | expected_fp8_stride = (d, num_tokens * d, 1) |
| 478 | assert fused_fp8.stride() == expected_fp8_stride, ( |
| 479 | f"FP8 stride mismatch: got {fused_fp8.stride()}, expected {expected_fp8_stride}" |
| 480 | ) |
| 481 | |
| 482 | # Scale: MN-major TMA-aligned layout. After fp8_einsum permutes |
| 483 | # [T,G,S] -> [G,T,S], T-dim should have stride 1. |
| 484 | # Our output is [T,G,S] = transpose of [G,T,S]. |
| 485 | # So fused_scale.permute(1,0,2) should have T-stride=1. |
| 486 | perm = fused_scale.permute(1, 0, 2) # [G, T, S] |
| 487 | assert perm.stride(1) == 1 or num_tokens == 1, ( |
| 488 | f"Scale T-stride (after permute to [G,T,S]) should be 1, got {perm.stride(1)}" |
| 489 | ) |
| 490 | |
| 491 | |
| 492 | @pytest.mark.parametrize("num_tokens", [1, 7, 32, 128]) |
| 493 | @torch.inference_mode() |
| 494 | def test_per_group_contiguity(num_tokens): |
| 495 | """FP8 per-group slices must be contiguous. Scale per-group slices |
| 496 | are column-major (T-stride=1) — not row-major contiguous, which is |
| 497 | correct for TMA loads.""" |
| 498 | num_heads, n_groups = 64, 8 |
| 499 | heads_per_group = num_heads // n_groups |
| 500 | max_pos = 4096 |
| 501 | device = "cuda" |
| 502 | |
| 503 | o = torch.randn( |
| 504 | num_tokens, num_heads, HEAD_DIM, device=device, dtype=torch.bfloat16 |
| 505 | ) |
| 506 | positions = torch.randint( |
| 507 | 0, max_pos, (num_tokens,), device=device, dtype=torch.long |
| 508 | ) |
| 509 | cos_sin_cache = make_cos_sin_cache(max_pos, device=device) |
| 510 | |
| 511 | fused_fp8, fused_scale = fused_inv_rope_fp8_quant( |
| 512 | o.clone(), |
| 513 | positions, |
| 514 | cos_sin_cache, |
| 515 | n_groups, |
| 516 | heads_per_group, |
| 517 | ) |
| 518 | |
| 519 | for g in range(n_groups): |
| 520 | fp8_slice = fused_fp8[:, g, :] |
| 521 | assert fp8_slice.is_contiguous(), ( |
| 522 | f"o_fp8[:, {g}, :] is not contiguous: " |
| 523 | f"shape={list(fp8_slice.shape)}, stride={list(fp8_slice.stride())}" |
| 524 | ) |
| 525 | |
| 526 | |
| 527 | @torch.inference_mode() |
| 528 | def test_scales_are_power_of_two(): |
| 529 | """Verify all scales are exact powers of 2 (UE8M0 property).""" |
| 530 | num_tokens, num_heads, n_groups = 32, 64, 8 |
| 531 | heads_per_group = num_heads // n_groups |
| 532 | max_pos = 4096 |
| 533 | device = "cuda" |
| 534 | |
| 535 | o = torch.randn( |
| 536 | num_tokens, num_heads, HEAD_DIM, device=device, dtype=torch.bfloat16 |
| 537 | ) |
| 538 | positions = torch.randint( |
| 539 | 0, max_pos, (num_tokens,), device=device, dtype=torch.long |
| 540 | ) |
| 541 | cos_sin_cache = make_cos_sin_cache(max_pos, device=device) |
| 542 | |
| 543 | _, fused_scale = fused_inv_rope_fp8_quant( |
| 544 | o.clone(), |
| 545 | positions, |
| 546 | cos_sin_cache, |
| 547 | n_groups, |
| 548 | heads_per_group, |
| 549 | ) |
| 550 | |
| 551 | # log2 of a power-of-two is an exact integer |
| 552 | log2_scales = torch.log2(fused_scale) |
| 553 | residual = (log2_scales - log2_scales.round()).abs() |
| 554 | assert residual.max() < 1e-5, ( |
| 555 | f"Not all scales are powers of 2: max log2 residual = {residual.max().item()}" |
| 556 | ) |
| 557 | |
| 558 | |
| 559 | @torch.inference_mode() |
| 560 | def test_nope_dims_unchanged(): |
| 561 | """Nope dimensions (first 448 per head) should only be quantized, |
| 562 | not rotated. Verify by dequantizing and comparing against |
| 563 | quantize-only reference (no RoPE).""" |
| 564 | num_tokens, num_heads, n_groups = 16, 64, 8 |
| 565 | heads_per_group = num_heads // n_groups |
| 566 | max_pos = 4096 |
| 567 | device = "cuda" |
| 568 | torch.manual_seed(0) |
| 569 | |
| 570 | o = torch.randn( |
| 571 | num_tokens, num_heads, HEAD_DIM, device=device, dtype=torch.bfloat16 |
| 572 | ) |
| 573 | positions = torch.randint( |
| 574 | 0, max_pos, (num_tokens,), device=device, dtype=torch.long |
| 575 | ) |
| 576 | cos_sin_cache = make_cos_sin_cache(max_pos, device=device) |
| 577 | |
| 578 | # Fused kernel result |
| 579 | fused_fp8, fused_scale = fused_inv_rope_fp8_quant( |
| 580 | o.clone(), |
| 581 | positions, |
| 582 | cos_sin_cache, |
| 583 | n_groups, |
| 584 | heads_per_group, |
| 585 | ) |
| 586 | |
| 587 | # Reference: quantize without RoPE (identity rotation) |
| 588 | # Create a zero-sin cache so RoPE is identity |
| 589 | zero_cache = torch.zeros_like(cos_sin_cache) |
| 590 | half = ROPE_DIM // 2 |
| 591 | zero_cache[:, :half] = 1.0 # cos = 1 |
| 592 | # sin = 0 (already zero) |
| 593 | |
| 594 | norope_fp8, norope_scale = fused_inv_rope_fp8_quant( |
| 595 | o.clone(), |
| 596 | positions, |
| 597 | zero_cache, |
| 598 | n_groups, |
| 599 | heads_per_group, |
| 600 | ) |
| 601 | |
| 602 | # Extract nope quant blocks only (first 3 of every 4 blocks per head) |
| 603 | chunks_per_head = HEAD_DIM // QUANT_GROUP_SIZE # 4 |
| 604 | |
| 605 | for h in range(heads_per_group): |
| 606 | for c in range(chunks_per_head - 1): # skip last chunk (has rope) |
| 607 | qb = h * chunks_per_head + c |
| 608 | start = qb * QUANT_GROUP_SIZE |
| 609 | end = start + QUANT_GROUP_SIZE |
| 610 | |
| 611 | fused_nope = fused_fp8[:, :, start:end].view(torch.uint8) |
| 612 | norope_nope = norope_fp8[:, :, start:end].view(torch.uint8) |
| 613 | assert torch.equal(fused_nope, norope_nope), ( |
| 614 | f"Nope block (head={h}, chunk={c}) differs between " |
| 615 | f"fused and no-rope reference" |
| 616 | ) |
| 617 | |
| 618 | fused_s = fused_scale[:, :, qb] |
| 619 | norope_s = norope_scale[:, :, qb] |
| 620 | assert torch.equal(fused_s, norope_s), ( |
| 621 | f"Nope scale (head={h}, chunk={c}) differs" |
| 622 | ) |
| 623 | |
| 624 | |
| 625 | @torch.inference_mode() |
| 626 | def test_single_token(): |
| 627 | """Edge case: single token.""" |
| 628 | num_tokens, num_heads, n_groups = 1, 64, 8 |
| 629 | heads_per_group = num_heads // n_groups |
| 630 | max_pos = 4096 |
| 631 | device = "cuda" |
| 632 | |
| 633 | o = torch.randn( |
| 634 | num_tokens, num_heads, HEAD_DIM, device=device, dtype=torch.bfloat16 |
| 635 | ) |
| 636 | positions = torch.tensor([42], device=device, dtype=torch.long) |
| 637 | cos_sin_cache = make_cos_sin_cache(max_pos, device=device) |
| 638 | |
| 639 | ref_fp8, ref_scale = reference_inv_rope_fp8_quant( |
| 640 | o.clone(), |
| 641 | positions, |
| 642 | cos_sin_cache, |
| 643 | n_groups, |
| 644 | heads_per_group, |
| 645 | ) |
| 646 | fused_fp8, fused_scale = fused_inv_rope_fp8_quant( |
| 647 | o.clone(), |
| 648 | positions, |
| 649 | cos_sin_cache, |
| 650 | n_groups, |
| 651 | heads_per_group, |
| 652 | ) |
| 653 | |
| 654 | assert_dequant_close(ref_fp8, ref_scale, fused_fp8, fused_scale) |
| 655 | |
| 656 | |
| 657 | @torch.inference_mode() |
| 658 | def test_zero_positions(): |
| 659 | """Edge case: all positions are 0.""" |
| 660 | num_tokens, num_heads, n_groups = 16, 64, 8 |
| 661 | heads_per_group = num_heads // n_groups |
| 662 | max_pos = 4096 |
| 663 | device = "cuda" |
| 664 | |
| 665 | o = torch.randn( |
| 666 | num_tokens, num_heads, HEAD_DIM, device=device, dtype=torch.bfloat16 |
| 667 | ) |
| 668 | positions = torch.zeros(num_tokens, device=device, dtype=torch.long) |
| 669 | cos_sin_cache = make_cos_sin_cache(max_pos, device=device) |
| 670 | |
| 671 | ref_fp8, ref_scale = reference_inv_rope_fp8_quant( |
| 672 | o.clone(), |
| 673 | positions, |
| 674 | cos_sin_cache, |
| 675 | n_groups, |
| 676 | heads_per_group, |
| 677 | ) |
| 678 | fused_fp8, fused_scale = fused_inv_rope_fp8_quant( |
| 679 | o.clone(), |
| 680 | positions, |
| 681 | cos_sin_cache, |
| 682 | n_groups, |
| 683 | heads_per_group, |
| 684 | ) |
| 685 | |
| 686 | assert_dequant_close(ref_fp8, ref_scale, fused_fp8, fused_scale) |
| 687 | |
| 688 | |
| 689 | @torch.inference_mode() |
| 690 | def test_large_values(): |
| 691 | """Edge case: values near FP8 saturation to test clamping.""" |
| 692 | num_tokens, num_heads, n_groups = 8, 64, 8 |
| 693 | heads_per_group = num_heads // n_groups |
| 694 | max_pos = 4096 |
| 695 | device = "cuda" |
| 696 | |
| 697 | # Create inputs with large values that will saturate FP8 |
| 698 | o = torch.randn( |
| 699 | num_tokens, num_heads, HEAD_DIM, device=device, dtype=torch.bfloat16 |
| 700 | ) |
| 701 | o = o * 1000.0 # scale up to force saturation |
| 702 | positions = torch.randint( |
| 703 | 0, max_pos, (num_tokens,), device=device, dtype=torch.long |
| 704 | ) |
| 705 | cos_sin_cache = make_cos_sin_cache(max_pos, device=device) |
| 706 | |
| 707 | ref_fp8, ref_scale = reference_inv_rope_fp8_quant( |
| 708 | o.clone(), |
| 709 | positions, |
| 710 | cos_sin_cache, |
| 711 | n_groups, |
| 712 | heads_per_group, |
| 713 | ) |
| 714 | fused_fp8, fused_scale = fused_inv_rope_fp8_quant( |
| 715 | o.clone(), |
| 716 | positions, |
| 717 | cos_sin_cache, |
| 718 | n_groups, |
| 719 | heads_per_group, |
| 720 | ) |
| 721 | |
| 722 | assert_dequant_close(ref_fp8, ref_scale, fused_fp8, fused_scale) |
| 723 | |
| 724 | |
| 725 | @torch.inference_mode() |
| 726 | def test_dequant_numerical_accuracy(): |
| 727 | """Verify dequantized values are close to the original (after inv RoPE).""" |
| 728 | num_tokens, num_heads, n_groups = 32, 64, 8 |
| 729 | heads_per_group = num_heads // n_groups |
| 730 | max_pos = 4096 |
| 731 | device = "cuda" |
| 732 | torch.manual_seed(0) |
| 733 | |
| 734 | o = torch.randn( |
| 735 | num_tokens, num_heads, HEAD_DIM, device=device, dtype=torch.bfloat16 |
| 736 | ) |
| 737 | positions = torch.randint( |
| 738 | 0, max_pos, (num_tokens,), device=device, dtype=torch.long |
| 739 | ) |
| 740 | cos_sin_cache = make_cos_sin_cache(max_pos, device=device) |
| 741 | |
| 742 | # Get the post-inv-RoPE values (ground truth before quantization) |
| 743 | o_after_rope = reference_inv_rope(o.clone(), positions, cos_sin_cache) |
| 744 | d = heads_per_group * HEAD_DIM |
| 745 | o_after_rope = o_after_rope.view(num_tokens, n_groups, d) |
| 746 | |
| 747 | # Get fused quantized output |
| 748 | fused_fp8, fused_scale = fused_inv_rope_fp8_quant( |
| 749 | o.clone(), |
| 750 | positions, |
| 751 | cos_sin_cache, |
| 752 | n_groups, |
| 753 | heads_per_group, |
| 754 | ) |
| 755 | |
| 756 | # Dequantize: broadcast scale [T, G, S] to [T, G, D] via repeat |
| 757 | S = d // QUANT_GROUP_SIZE |
| 758 | scale_expanded = ( |
| 759 | fused_scale.unsqueeze(-1) |
| 760 | .expand(num_tokens, n_groups, S, QUANT_GROUP_SIZE) |
| 761 | .reshape(num_tokens, n_groups, d) |
| 762 | ) |
| 763 | dequant = fused_fp8.float() * scale_expanded |
| 764 | |
| 765 | # Check relative error. |
| 766 | # FP8 e4m3 with UE8M0 (power-of-two scales that round UP) quantizes more |
| 767 | # coarsely than optimal scaling. Both paths rotate in fp32, so the bulk |
| 768 | # of the error comes from UE8M0 quantization itself (~10-12% typical). |
| 769 | o_gt = o_after_rope.transpose(0, 1).contiguous().transpose(0, 1) |
| 770 | dequant_contig = dequant.transpose(0, 1).contiguous().transpose(0, 1) |
| 771 | |
| 772 | abs_err = (dequant_contig.float() - o_gt.float()).abs() |
| 773 | rel_err = abs_err / (o_gt.float().abs().clamp(min=1e-6)) |
| 774 | mean_rel_err = rel_err.mean().item() |
| 775 | |
| 776 | assert mean_rel_err < 0.15, ( |
| 777 | f"Mean relative error too high: {mean_rel_err:.4f} (expected < 0.15)" |
| 778 | ) |
| 779 | |
| 780 | |
| 781 | def _unfused_inv_rope_fp8_quant( |
| 782 | o: torch.Tensor, |
| 783 | positions: torch.Tensor, |
| 784 | cos_sin_cache: torch.Tensor, |
| 785 | n_groups: int, |
| 786 | heads_per_group: int, |
| 787 | nope_dim: int = NOPE_DIM, |
| 788 | rope_dim: int = ROPE_DIM, |
| 789 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 790 | """Unfused path matching deepseek_v4/attention.py:295-310. |
| 791 | |
| 792 | Uses the production CUDA RoPE kernel + per_token_group_quant_fp8. |
| 793 | """ |
| 794 | from vllm import _custom_ops as ops |
| 795 | from vllm.model_executor.layers.quantization.utils.fp8_utils import ( |
| 796 | per_token_group_quant_fp8, |
| 797 | ) |
| 798 | |
| 799 | head_dim = o.shape[-1] |
| 800 | rope_dim_offset = head_dim - rope_dim |
| 801 | |
| 802 | # Step 1: In-place CUDA RoPE (same as production) |
| 803 | ops.rotary_embedding( |
| 804 | positions, |
| 805 | o, |
| 806 | None, |
| 807 | head_dim, |
| 808 | cos_sin_cache, |
| 809 | False, # is_neox=False for DeepseekV4 (GPT-J style) |
| 810 | rope_dim_offset=rope_dim_offset, |
| 811 | inverse=True, |
| 812 | ) |
| 813 | |
| 814 | # Step 2: Reshape + quant + reshape (same as production) |
| 815 | T = o.shape[0] |
| 816 | d = heads_per_group * head_dim |
| 817 | o = o.view(T, n_groups, -1) |
| 818 | o_flat = o.transpose(0, 1).contiguous().reshape(-1, d) |
| 819 | o_fp8, o_scale = per_token_group_quant_fp8( |
| 820 | o_flat, |
| 821 | group_size=QUANT_GROUP_SIZE, |
| 822 | use_ue8m0=True, |
| 823 | ) |
| 824 | o_fp8 = o_fp8.view(n_groups, T, d).transpose(0, 1) |
| 825 | o_scale = o_scale.view(n_groups, T, -1).transpose(0, 1) |
| 826 | return o_fp8, o_scale |
| 827 | |
| 828 | |
| 829 | # ========================================================================= |
| 830 | # End-to-end test including fp8_einsum |
| 831 | # ========================================================================= |
| 832 | |
| 833 | |
| 834 | @pytest.mark.parametrize("num_tokens", [1, 7, 32, 128, 1024]) |
| 835 | @pytest.mark.parametrize( |
| 836 | "num_heads,n_groups", |
| 837 | [(64, 8)], |
| 838 | ids=["H64_G8"], |
| 839 | ) |
| 840 | @torch.inference_mode() |
| 841 | def test_einsum_end_to_end(num_tokens, num_heads, n_groups): |
| 842 | """End-to-end: fused inv_rope+quant → fp8_einsum must match |
| 843 | unfused CUDA_rope+quant → fp8_einsum bitwise. |
| 844 | |
| 845 | This catches stride/layout bugs that only manifest when the einsum |
| 846 | kernel actually consumes the quantized activations. |
| 847 | """ |
| 848 | from vllm.models.deepseek_v4.nvidia.ops.o_proj import deep_gemm_fp8_o_proj |
| 849 | from vllm.utils.deep_gemm import ( |
| 850 | fp8_einsum, |
| 851 | is_deep_gemm_supported, |
| 852 | per_block_cast_to_fp8, |
| 853 | transform_sf_into_required_layout, |
| 854 | ) |
| 855 | |
| 856 | if not is_deep_gemm_supported(): |
| 857 | pytest.skip("DeepGEMM not supported on this platform") |
| 858 | |
| 859 | def ceil_div(a: int, b: int) -> int: |
| 860 | return (a + b - 1) // b |
| 861 | |
| 862 | heads_per_group = num_heads // n_groups |
| 863 | d = heads_per_group * HEAD_DIM |
| 864 | o_lora_rank = 1024 |
| 865 | max_pos = 4096 |
| 866 | device = "cuda" |
| 867 | torch.manual_seed(0) |
| 868 | |
| 869 | o = torch.randn( |
| 870 | num_tokens, num_heads, HEAD_DIM, device=device, dtype=torch.bfloat16 |
| 871 | ) |
| 872 | positions = torch.randint( |
| 873 | 0, max_pos, (num_tokens,), device=device, dtype=torch.long |
| 874 | ) |
| 875 | cos_sin_cache = make_cos_sin_cache(max_pos, device=device) |
| 876 | |
| 877 | # -- Weight quantization (shared between both paths) -- |
| 878 | w = torch.randn(n_groups, o_lora_rank, d, device=device, dtype=torch.bfloat16) |
| 879 | w_fp8 = torch.empty_like(w, dtype=torch.float8_e4m3fn) |
| 880 | w_scale = torch.empty( |
| 881 | n_groups, |
| 882 | ceil_div(o_lora_rank, 128), |
| 883 | ceil_div(d, 128), |
| 884 | device=device, |
| 885 | dtype=torch.float32, |
| 886 | ) |
| 887 | for g in range(n_groups): |
| 888 | w_fp8[g], w_scale[g] = per_block_cast_to_fp8(w[g], use_ue8m0=True) |
| 889 | |
| 890 | recipe = (1, 1, 128) |
| 891 | w_scale_t = transform_sf_into_required_layout( |
| 892 | sf=w_scale, |
| 893 | mn=o_lora_rank, |
| 894 | k=d, |
| 895 | recipe=(1, 128, 128), |
| 896 | num_groups=n_groups, |
| 897 | is_sfa=False, |
| 898 | ) |
| 899 | |
| 900 | # -- UNFUSED path -- |
| 901 | ref_fp8, ref_scale = _unfused_inv_rope_fp8_quant( |
| 902 | o.clone(), |
| 903 | positions, |
| 904 | cos_sin_cache, |
| 905 | n_groups, |
| 906 | heads_per_group, |
| 907 | ) |
| 908 | z_ref = torch.empty( |
| 909 | num_tokens, n_groups, o_lora_rank, device=device, dtype=torch.bfloat16 |
| 910 | ) |
| 911 | fp8_einsum( |
| 912 | "bhr,hdr->bhd", (ref_fp8, ref_scale), (w_fp8, w_scale_t), z_ref, recipe=recipe |
| 913 | ) |
| 914 | |
| 915 | # -- FUSED path -- |
| 916 | wo_a = torch.nn.Module() |
| 917 | wo_a.register_buffer("weight", w_fp8) |
| 918 | wo_a.register_buffer("weight_scale", w_scale_t) |
| 919 | z_fused = deep_gemm_fp8_o_proj( |
| 920 | o.clone(), |
| 921 | positions, |
| 922 | cos_sin_cache, |
| 923 | wo_a, |
| 924 | torch.nn.Identity(), |
| 925 | n_groups=n_groups, |
| 926 | heads_per_group=heads_per_group, |
| 927 | nope_dim=NOPE_DIM, |
| 928 | rope_dim=ROPE_DIM, |
| 929 | o_lora_rank=o_lora_rank, |
| 930 | einsum_recipe=recipe, |
| 931 | tma_aligned_scales=True, |
| 932 | ).view_as(z_ref) |
| 933 | |
| 934 | # -- Checks -- |
| 935 | # Einsum output: Triton and CUDA both rotate in fp32 now, so diffs |
| 936 | # come from fp32 ordering and UE8M0 boundary shifts only. |
| 937 | # Use relative diff (same metric as deep_gemm.testing.calc_diff). |
| 938 | def calc_diff(x, y): |
| 939 | x, y = x.double(), y.double() |
| 940 | denominator = (x * x + y * y).sum() |
| 941 | sim = 2 * (x * y).sum() / denominator |
| 942 | return 1 - sim |
| 943 | |
| 944 | z_diff = calc_diff(z_fused, z_ref) |
| 945 | assert z_diff < 0.01, ( |
| 946 | f"Einsum output diff too large: {z_diff:.6f} (expected < 0.01)" |
| 947 | ) |
| 948 | |
| 949 | |
| 950 | @pytest.mark.parametrize("num_tokens", [1, 32, 256]) |
| 951 | @torch.inference_mode() |
| 952 | def test_with_real_deepseek_v4_rope(num_tokens, default_vllm_config): |
| 953 | """Test with real DeepseekV4ScalingRotaryEmbedding (GPT-J style, |
| 954 | mscale=0, YaRN scaling) matching the production config.""" |
| 955 | num_heads = 64 |
| 956 | n_groups = 8 |
| 957 | heads_per_group = num_heads // n_groups |
| 958 | device = "cuda" |
| 959 | torch.manual_seed(0) |
| 960 | |
| 961 | # Build YaRN-scaled cos_sin_cache matching real DeepSeek V3/V4 config |
| 962 | # (mscale=0 → mscale=1.0, so no magnitude scaling) |
| 963 | from vllm.model_executor.layers.rotary_embedding.common import ( |
| 964 | yarn_find_correction_range, |
| 965 | yarn_linear_ramp_mask, |
| 966 | ) |
| 967 | |
| 968 | scaling_factor = 16 |
| 969 | base = 10000.0 |
| 970 | max_pos = 65536 |
| 971 | beta_fast, beta_slow = 32, 1 |
| 972 | |
| 973 | pos_freqs = base ** ( |
| 974 | torch.arange(0, ROPE_DIM, 2, dtype=torch.float32, device=device) / ROPE_DIM |
| 975 | ) |
| 976 | inv_freq_extra = 1.0 / pos_freqs |
| 977 | inv_freq_interp = 1.0 / (scaling_factor * pos_freqs) |
| 978 | low, high = yarn_find_correction_range( |
| 979 | beta_fast, beta_slow, ROPE_DIM, base, max_pos |
| 980 | ) |
| 981 | mask = 1 - yarn_linear_ramp_mask(low, high, ROPE_DIM // 2, dtype=torch.float32).to( |
| 982 | device |
| 983 | ) |
| 984 | inv_freq = inv_freq_interp * (1 - mask) + inv_freq_extra * mask |
| 985 | t = torch.arange(max_pos * scaling_factor, device=device, dtype=torch.float32) |
| 986 | freqs = torch.outer(t, inv_freq) |
| 987 | # mscale=0 → yarn_get_mscale returns 1.0 |
| 988 | cos_sin_cache = torch.cat([freqs.cos(), freqs.sin()], dim=-1) # fp32 |
| 989 | |
| 990 | o = torch.randn( |
| 991 | num_tokens, num_heads, HEAD_DIM, device=device, dtype=torch.bfloat16 |
| 992 | ) |
| 993 | positions = torch.randint(0, 4096, (num_tokens,), device=device, dtype=torch.long) |
| 994 | |
| 995 | # UNFUSED: CUDA RoPE with is_neox=False (GPT-J) |
| 996 | from vllm import _custom_ops as ops |
| 997 | from vllm.model_executor.layers.quantization.utils.fp8_utils import ( |
| 998 | per_token_group_quant_fp8, |
| 999 | ) |
| 1000 | |
| 1001 | o_unfused = o.clone() |
| 1002 | ops.rotary_embedding( |
| 1003 | positions, |
| 1004 | o_unfused, |
| 1005 | None, |
| 1006 | HEAD_DIM, |
| 1007 | cos_sin_cache, |
| 1008 | False, # is_neox=False (GPT-J style) |
| 1009 | rope_dim_offset=NOPE_DIM, |
| 1010 | inverse=True, |
| 1011 | ) |
| 1012 | d = heads_per_group * HEAD_DIM |
| 1013 | T = num_tokens |
| 1014 | o_unfused = o_unfused.view(T, n_groups, d) |
| 1015 | o_flat = o_unfused.transpose(0, 1).contiguous().reshape(-1, d) |
| 1016 | ref_fp8, ref_scale = per_token_group_quant_fp8( |
| 1017 | o_flat, |
| 1018 | group_size=QUANT_GROUP_SIZE, |
| 1019 | use_ue8m0=True, |
| 1020 | ) |
| 1021 | ref_fp8 = ref_fp8.view(n_groups, T, d).transpose(0, 1) |
| 1022 | ref_scale = ref_scale.view(n_groups, T, -1).transpose(0, 1) |
| 1023 | |
| 1024 | # FUSED: use the real YaRN-scaled cos_sin_cache |
| 1025 | fused_fp8, fused_scale = fused_inv_rope_fp8_quant( |
| 1026 | o.clone(), |
| 1027 | positions, |
| 1028 | cos_sin_cache, |
| 1029 | n_groups, |
| 1030 | heads_per_group, |
| 1031 | ) |
| 1032 | |
| 1033 | # Scales must match exactly (same UE8M0 algorithm) |
| 1034 | # Compare via dequant (Triton bf16 rotation may differ from CUDA by 1 ULP) |
| 1035 | assert_dequant_close( |
| 1036 | ref_fp8, ref_scale, fused_fp8, fused_scale, msg="Real DeepSeek V4 rope" |
| 1037 | ) |