| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Unit tests for the horizontally-fused deepseek_v32 (NVIDIA SM100) Triton |
| 4 | kernels used by the specialized DSA model: |
| 5 | |
| 6 | fused_norm_rope |
| 7 | - q : q_lora RMSNorm |
| 8 | - kv : kv_lora RMSNorm + (interleaved) RoPE on k_pe + MLA cache insert |
| 9 | (bf16 or per-tensor fp8) |
| 10 | - idx: indexer-K LayerNorm + RoPE (interleaved or NeoX) + UE8M0 fp8 quant + |
| 11 | packed indexer cache insert; plus the top-k buffer (-1) fill |
| 12 | fused_q |
| 13 | - mqa: ql_nope + (interleaved) RoPE'd q_pe, concat-quantized to the fp8 MQA |
| 14 | query |
| 15 | - idx: indexer-Q RoPE (interleaved or NeoX) + UE8M0 fp8 quant + folded |
| 16 | index weights |
| 17 | fused_eh_norm (MTP): zero-at-pos-0 + enorm RMSNorm(embeds) + hnorm |
| 18 | RMSNorm(prev), concatenated side-by-side |
| 19 | |
| 20 | Each kernel is compared against a PyTorch reference. The kernel keeps the whole |
| 21 | pipeline in fp32 and rounds once, so it can land on the opposite side of a |
| 22 | round-to-nearest tie from the reference for a few elements: deterministic fp8 |
| 23 | outputs are checked within 1 representable-step (ULP); bf16 norm/RoPE outputs use |
| 24 | rtol/atol=1e-2 (the tolerance the sibling deepseek_v4 fused-kernel test uses). |
| 25 | """ |
| 26 | |
| 27 | from functools import partial |
| 28 | from typing import NamedTuple |
| 29 | |
| 30 | import pytest |
| 31 | import torch |
| 32 | |
| 33 | from vllm.models.deepseek_v32.common import kernels as K |
| 34 | from vllm.platforms import current_platform |
| 35 | |
| 36 | FP8 = torch.float8_e4m3fn |
| 37 | FP8_MAX = 448.0 |
| 38 | |
| 39 | EPS = 1e-6 |
| 40 | |
| 41 | |
| 42 | class ModelConfig(NamedTuple): |
| 43 | """Per-model attention dims (num_heads is the TP8 local count).""" |
| 44 | |
| 45 | q_lora: int |
| 46 | kv_lora: int |
| 47 | rope_dim: int |
| 48 | num_heads: int |
| 49 | index_heads: int |
| 50 | index_head_dim: int |
| 51 | hidden: int |
| 52 | |
| 53 | |
| 54 | MODEL_CONFIGS = [ |
| 55 | ModelConfig(2048, 512, 64, 8, 32, 128, 6144), # GLM-5.2 |
| 56 | ModelConfig(1536, 512, 64, 16, 64, 128, 7168), # DeepSeek-V3.2 |
| 57 | ] |
| 58 | MODEL_IDS = ["glm5.2", "dsv3.2"] |
| 59 | |
| 60 | pytestmark = pytest.mark.skipif( |
| 61 | not current_platform.supports_fp8(), |
| 62 | reason="deepseek_v32 fused kernels require fp8 support", |
| 63 | ) |
| 64 | |
| 65 | |
| 66 | # ── reference helpers ──────────────────────────────────────────────────────── |
| 67 | |
| 68 | |
| 69 | def make_cos_sin(max_pos: int, rot_dim: int, device) -> torch.Tensor: |
| 70 | """cos||sin cache: row[pos] = [cos(theta)(rot/2), sin(theta)(rot/2)].""" |
| 71 | half = rot_dim // 2 |
| 72 | inv_freq = 1.0 / ( |
| 73 | 10000.0 ** (torch.arange(0, half, dtype=torch.float32, device=device) / half) |
| 74 | ) |
| 75 | t = torch.arange(max_pos, dtype=torch.float32, device=device) |
| 76 | freqs = torch.einsum("i,j->ij", t, inv_freq) |
| 77 | return torch.cat([freqs.cos(), freqs.sin()], dim=-1) |
| 78 | |
| 79 | |
| 80 | def rms_norm(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: |
| 81 | """RMSNorm matching kernels._rms_norm (fp32, eps inside rsqrt). Returns fp32.""" |
| 82 | xf = x.float() |
| 83 | ms = xf.pow(2).mean(dim=-1, keepdim=True) |
| 84 | return xf * torch.rsqrt(ms + EPS) * w.float() |
| 85 | |
| 86 | |
| 87 | def layer_norm(x: torch.Tensor, w: torch.Tensor, b: torch.Tensor) -> torch.Tensor: |
| 88 | xf = x.float() |
| 89 | mean = xf.mean(dim=-1, keepdim=True) |
| 90 | var = (xf - mean).pow(2).mean(dim=-1, keepdim=True) |
| 91 | return (xf - mean) * torch.rsqrt(var + EPS) * w.float() + b.float() |
| 92 | |
| 93 | |
| 94 | def rope( |
| 95 | x: torch.Tensor, pos: torch.Tensor, cos_sin: torch.Tensor, interleave: bool |
| 96 | ) -> torch.Tensor: |
| 97 | """Apply RoPE to the first ``rot_dim`` elements of x's last dim. |
| 98 | |
| 99 | x: [..., head_dim] fp32. ``cos_sin`` is [max_pos, rot_dim]. ``interleave`` |
| 100 | selects adjacent-pair (GLM) vs split-half NeoX (DeepSeek-V3.2) layout. |
| 101 | """ |
| 102 | rot = cos_sin.shape[-1] |
| 103 | half = rot // 2 |
| 104 | cs = cos_sin[pos.long()] |
| 105 | cos, sin = cs[..., :half], cs[..., half:] |
| 106 | out = x.float().clone() |
| 107 | r = out[..., :rot] |
| 108 | if interleave: |
| 109 | x1, x2 = r[..., 0::2].clone(), r[..., 1::2].clone() |
| 110 | r[..., 0::2] = x1 * cos - x2 * sin |
| 111 | r[..., 1::2] = x2 * cos + x1 * sin |
| 112 | else: |
| 113 | x1, x2 = r[..., :half].clone(), r[..., half:].clone() |
| 114 | r[..., :half] = x1 * cos - x2 * sin |
| 115 | r[..., half:] = x2 * cos + x1 * sin |
| 116 | return out |
| 117 | |
| 118 | |
| 119 | def ue8m0_quant(vals: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| 120 | """Per-row (last dim) UE8M0 fp8 quant matching kernels._fp8_ue8m0_quantize.""" |
| 121 | amax = vals.float().abs().amax(dim=-1, keepdim=True) |
| 122 | scale = torch.clamp(amax, min=1e-4) / FP8_MAX |
| 123 | scale = torch.exp2(torch.ceil(torch.log2(scale))) |
| 124 | q = (vals.float() / scale).to(FP8) |
| 125 | return q, scale.squeeze(-1) |
| 126 | |
| 127 | |
| 128 | def _bf16_ulp(a: torch.Tensor, b: torch.Tensor) -> int: |
| 129 | def key(t): |
| 130 | u = t.contiguous().view(torch.int16).to(torch.int64) & 0xFFFF |
| 131 | return torch.where(u >= 0x8000, 0xFFFF - u, u + 0x8000) |
| 132 | |
| 133 | return int((key(a) - key(b)).abs().max().item()) |
| 134 | |
| 135 | |
| 136 | def _fp8_ulp(a: torch.Tensor, b: torch.Tensor) -> int: |
| 137 | def key(t): |
| 138 | u = t.contiguous().view(torch.uint8).to(torch.int64) |
| 139 | return torch.where(u >= 0x80, 0xFF - u, u + 0x80) |
| 140 | |
| 141 | return int((key(a) - key(b)).abs().max().item()) |
| 142 | |
| 143 | |
| 144 | def assert_bf16(got: torch.Tensor, ref_fp32: torch.Tensor, msg: str): |
| 145 | # Kernel keeps RMSNorm/RoPE in fp32 and rounds to bf16 once; the fp32 |
| 146 | # reduction/FMA order differs from torch, so a few elements land on the |
| 147 | # opposite side of a round-to-nearest tie. Use the same tolerance the |
| 148 | # sibling deepseek_v4 fused-kernel test uses for this bf16 norm+rope class. |
| 149 | torch.testing.assert_close( |
| 150 | got.float(), ref_fp32.float(), rtol=1e-2, atol=1e-2, msg=lambda m: f"{msg}: {m}" |
| 151 | ) |
| 152 | |
| 153 | |
| 154 | def assert_fp8(got: torch.Tensor, ref: torch.Tensor, msg: str): |
| 155 | assert _fp8_ulp(got, ref) <= 1, f"{msg}: >1 fp8 ULP" |
| 156 | |
| 157 | |
| 158 | # ── fused_norm_rope ────────────────────────────────────────────────────────── |
| 159 | |
| 160 | |
| 161 | @pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS) |
| 162 | @pytest.mark.parametrize("num_tokens", [1, 4, 17, 512, 4096]) |
| 163 | @pytest.mark.parametrize("index_interleave", [True, False]) |
| 164 | @pytest.mark.parametrize("mla_dtype", ["auto", "bfloat16", "fp8"]) |
| 165 | def test_fused_norm_rope( |
| 166 | num_tokens: int, index_interleave: bool, mla_dtype: str, cfg: ModelConfig |
| 167 | ): |
| 168 | torch.manual_seed(0) |
| 169 | dev = "cuda" |
| 170 | max_pos = 8192 |
| 171 | pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos |
| 172 | |
| 173 | q_c = torch.randn(num_tokens, cfg.q_lora, device=dev, dtype=torch.bfloat16) |
| 174 | kv_c = torch.randn(num_tokens, cfg.kv_lora, device=dev, dtype=torch.bfloat16) |
| 175 | k_pe = torch.randn(num_tokens, cfg.rope_dim, device=dev, dtype=torch.bfloat16) |
| 176 | qw = torch.randn(cfg.q_lora, device=dev, dtype=torch.bfloat16) |
| 177 | kvw = torch.randn(cfg.kv_lora, device=dev, dtype=torch.bfloat16) |
| 178 | ik = torch.randn(num_tokens, cfg.index_head_dim, device=dev, dtype=torch.bfloat16) |
| 179 | ikw = torch.randn(cfg.index_head_dim, device=dev, dtype=torch.float32) |
| 180 | ikb = torch.randn(cfg.index_head_dim, device=dev, dtype=torch.float32) |
| 181 | |
| 182 | mla_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev) # MLA k_pe: interleaved |
| 183 | idx_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev) |
| 184 | |
| 185 | bs = max_pos # single block covering all tokens |
| 186 | mla_dim = cfg.kv_lora + cfg.rope_dim |
| 187 | mla_fp8 = mla_dtype == "fp8" |
| 188 | if mla_fp8: |
| 189 | mla_cache = torch.zeros(1, bs, mla_dim, device=dev, dtype=torch.uint8) |
| 190 | mla_k_scale = torch.tensor([0.3], device=dev, dtype=torch.float32) |
| 191 | else: |
| 192 | mla_cache = torch.zeros(1, bs, mla_dim, device=dev, dtype=torch.bfloat16) |
| 193 | mla_k_scale = None |
| 194 | idx_row = cfg.index_head_dim + cfg.index_head_dim // 128 * 4 # 132 |
| 195 | idx_cache = torch.zeros(1, bs, idx_row, device=dev, dtype=torch.uint8) |
| 196 | slot = torch.arange(num_tokens, device=dev, dtype=torch.int64) |
| 197 | topk = torch.full((num_tokens, 2048), 7, device=dev, dtype=torch.int32) |
| 198 | |
| 199 | q_out = K.fused_norm_rope( |
| 200 | pos, |
| 201 | q_c, |
| 202 | qw, |
| 203 | EPS, |
| 204 | kv_c, |
| 205 | kvw, |
| 206 | EPS, |
| 207 | k_pe, |
| 208 | mla_cos_sin, |
| 209 | ik, |
| 210 | ikw, |
| 211 | ikb, |
| 212 | EPS, |
| 213 | idx_cos_sin, |
| 214 | topk, |
| 215 | slot_mapping=slot, |
| 216 | indexer_k_cache=idx_cache, |
| 217 | mla_kv_cache=mla_cache, |
| 218 | mla_kv_cache_dtype=mla_dtype, |
| 219 | mla_k_scale=mla_k_scale, |
| 220 | has_indexer=True, |
| 221 | index_rope_interleave=index_interleave, |
| 222 | ) |
| 223 | |
| 224 | # q_lora RMSNorm |
| 225 | assert_bf16(q_out, rms_norm(q_c, qw), "q_c rmsnorm") |
| 226 | |
| 227 | # MLA cache: [kv_c_normed | k_pe_roped(interleaved)] |
| 228 | kv_ref = rms_norm(kv_c, kvw) |
| 229 | kpe_ref = rope(k_pe.float(), pos, mla_cos_sin, interleave=True) |
| 230 | if mla_fp8: |
| 231 | cache = mla_cache.view(FP8)[0, :num_tokens] |
| 232 | s = mla_k_scale.item() |
| 233 | assert_fp8(cache[:, : cfg.kv_lora], (kv_ref / s).to(FP8), "MLA kv fp8") |
| 234 | assert_fp8(cache[:, cfg.kv_lora :], (kpe_ref / s).to(FP8), "MLA k_pe fp8") |
| 235 | else: |
| 236 | cache = mla_cache[0, :num_tokens] |
| 237 | assert_bf16(cache[:, : cfg.kv_lora], kv_ref, "MLA kv bf16") |
| 238 | assert_bf16(cache[:, cfg.kv_lora :], kpe_ref, "MLA k_pe bf16") |
| 239 | |
| 240 | # Indexer-K cache (packed [bs*head_dim fp8 | bs*4 fp32 scale]). |
| 241 | ik_ref = layer_norm(ik, ikw, ikb) |
| 242 | ik_ref = rope(ik_ref, pos, idx_cos_sin, interleave=index_interleave) |
| 243 | q_ref, s_ref = ue8m0_quant(ik_ref) |
| 244 | flat = idx_cache[0].reshape(-1) |
| 245 | vals = flat[: bs * cfg.index_head_dim].view(FP8).reshape(bs, cfg.index_head_dim) |
| 246 | scales = flat[bs * cfg.index_head_dim :].view(torch.float32) |
| 247 | assert_fp8(vals[:num_tokens], q_ref, "indexer-K fp8") |
| 248 | torch.testing.assert_close(scales[:num_tokens], s_ref, rtol=0, atol=0) |
| 249 | |
| 250 | # Top-k buffer cleared to -1 on indexer layers. |
| 251 | assert (topk == -1).all(), "topk buffer not cleared on indexer layer" |
| 252 | |
| 253 | |
| 254 | @pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS) |
| 255 | def test_fused_norm_rope_packed_indexer_block_stride(cfg: ModelConfig): |
| 256 | """Indexer writes use their HMA block stride and independent slot map.""" |
| 257 | torch.manual_seed(7) |
| 258 | dev = "cuda" |
| 259 | num_tokens = block_size = 2 |
| 260 | pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) |
| 261 | q_c = torch.randn(num_tokens, cfg.q_lora, device=dev, dtype=torch.bfloat16) |
| 262 | kv_c = torch.randn(num_tokens, cfg.kv_lora, device=dev, dtype=torch.bfloat16) |
| 263 | k_pe = torch.randn(num_tokens, cfg.rope_dim, device=dev, dtype=torch.bfloat16) |
| 264 | qw = torch.randn(cfg.q_lora, device=dev, dtype=torch.bfloat16) |
| 265 | kvw = torch.randn(cfg.kv_lora, device=dev, dtype=torch.bfloat16) |
| 266 | ik = torch.randn(num_tokens, cfg.index_head_dim, device=dev, dtype=torch.bfloat16) |
| 267 | ikw = torch.randn(cfg.index_head_dim, device=dev, dtype=torch.float32) |
| 268 | ikb = torch.randn(cfg.index_head_dim, device=dev, dtype=torch.float32) |
| 269 | cos_sin = make_cos_sin(32, cfg.rope_dim, dev) |
| 270 | |
| 271 | idx_row = cfg.index_head_dim + cfg.index_head_dim // 128 * 4 |
| 272 | packed_block_stride = block_size * idx_row + 64 |
| 273 | backing = torch.zeros(2 * packed_block_stride, device=dev, dtype=torch.uint8) |
| 274 | idx_cache = torch.as_strided( |
| 275 | backing, |
| 276 | (2, block_size, idx_row), |
| 277 | (packed_block_stride, idx_row, 1), |
| 278 | ) |
| 279 | mla_cache = torch.zeros( |
| 280 | 1, |
| 281 | block_size, |
| 282 | cfg.kv_lora + cfg.rope_dim, |
| 283 | device=dev, |
| 284 | dtype=torch.bfloat16, |
| 285 | ) |
| 286 | topk = torch.zeros(num_tokens, 8, device=dev, dtype=torch.int32) |
| 287 | |
| 288 | K.fused_norm_rope( |
| 289 | pos, |
| 290 | q_c, |
| 291 | qw, |
| 292 | EPS, |
| 293 | kv_c, |
| 294 | kvw, |
| 295 | EPS, |
| 296 | k_pe, |
| 297 | cos_sin, |
| 298 | ik, |
| 299 | ikw, |
| 300 | ikb, |
| 301 | EPS, |
| 302 | cos_sin, |
| 303 | topk, |
| 304 | slot_mapping=torch.arange(num_tokens, device=dev, dtype=torch.int64), |
| 305 | indexer_slot_mapping=torch.arange( |
| 306 | block_size, block_size + num_tokens, device=dev, dtype=torch.int64 |
| 307 | ), |
| 308 | indexer_k_cache=idx_cache, |
| 309 | mla_kv_cache=mla_cache, |
| 310 | has_indexer=True, |
| 311 | ) |
| 312 | |
| 313 | ik_ref = rope(layer_norm(ik, ikw, ikb), pos, cos_sin, interleave=False) |
| 314 | q_ref, s_ref = ue8m0_quant(ik_ref) |
| 315 | packed = idx_cache[1].reshape(-1) |
| 316 | values = ( |
| 317 | packed[: block_size * cfg.index_head_dim] |
| 318 | .view(FP8) |
| 319 | .view(block_size, cfg.index_head_dim) |
| 320 | ) |
| 321 | scales = packed[block_size * cfg.index_head_dim :].view(torch.float32) |
| 322 | assert_fp8(values, q_ref, "packed indexer-K fp8") |
| 323 | torch.testing.assert_close(scales, s_ref, rtol=0, atol=0) |
| 324 | assert (backing[block_size * idx_row : packed_block_stride] == 0).all() |
| 325 | |
| 326 | |
| 327 | @pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS) |
| 328 | def test_fused_norm_rope_materializes_qk_without_local_cache_slots( |
| 329 | cfg: ModelConfig, |
| 330 | ): |
| 331 | """DCP non-owner ranks still need valid query shards for query AllGather.""" |
| 332 | torch.manual_seed(7) |
| 333 | dev = "cuda" |
| 334 | num_tokens = 4 |
| 335 | max_pos = 16 |
| 336 | pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) |
| 337 | |
| 338 | q_c = torch.randn(num_tokens, cfg.q_lora, device=dev, dtype=torch.bfloat16) |
| 339 | kv_c = torch.randn(num_tokens, cfg.kv_lora, device=dev, dtype=torch.bfloat16) |
| 340 | k_pe = torch.randn(num_tokens, cfg.rope_dim, device=dev, dtype=torch.bfloat16) |
| 341 | qw = torch.randn(cfg.q_lora, device=dev, dtype=torch.bfloat16) |
| 342 | kvw = torch.randn(cfg.kv_lora, device=dev, dtype=torch.bfloat16) |
| 343 | ik = torch.randn(num_tokens, cfg.index_head_dim, device=dev, dtype=torch.bfloat16) |
| 344 | ikw = torch.randn(cfg.index_head_dim, device=dev, dtype=torch.float32) |
| 345 | ikb = torch.randn(cfg.index_head_dim, device=dev, dtype=torch.float32) |
| 346 | cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev) |
| 347 | |
| 348 | mla_cache = torch.zeros( |
| 349 | 1, max_pos, cfg.kv_lora + cfg.rope_dim, device=dev, dtype=torch.bfloat16 |
| 350 | ) |
| 351 | idx_row = cfg.index_head_dim + cfg.index_head_dim // 128 * 4 |
| 352 | idx_cache = torch.zeros(1, max_pos, idx_row, device=dev, dtype=torch.uint8) |
| 353 | no_local_slots = torch.full((num_tokens,), -1, device=dev, dtype=torch.int64) |
| 354 | topk = torch.full((num_tokens, 2048), 7, device=dev, dtype=torch.int32) |
| 355 | kv_out = torch.empty_like(kv_c) |
| 356 | kpe_out = torch.empty_like(k_pe) |
| 357 | |
| 358 | q_out = K.fused_norm_rope( |
| 359 | pos, |
| 360 | q_c, |
| 361 | qw, |
| 362 | EPS, |
| 363 | kv_c, |
| 364 | kvw, |
| 365 | EPS, |
| 366 | k_pe, |
| 367 | cos_sin, |
| 368 | ik, |
| 369 | ikw, |
| 370 | ikb, |
| 371 | EPS, |
| 372 | cos_sin, |
| 373 | topk, |
| 374 | slot_mapping=no_local_slots, |
| 375 | indexer_k_cache=idx_cache, |
| 376 | mla_kv_cache=mla_cache, |
| 377 | mla_kv_cache_dtype="auto", |
| 378 | mla_k_scale=None, |
| 379 | has_indexer=True, |
| 380 | index_rope_interleave=True, |
| 381 | kv_c_out=kv_out, |
| 382 | k_pe_out=kpe_out, |
| 383 | ) |
| 384 | |
| 385 | assert_bf16(q_out, rms_norm(q_c, qw), "q_c rmsnorm without local cache slots") |
| 386 | assert_bf16(kv_out, rms_norm(kv_c, kvw), "kv_c rmsnorm without local slots") |
| 387 | assert_bf16( |
| 388 | kpe_out, |
| 389 | rope(k_pe.float(), pos, cos_sin, interleave=True), |
| 390 | "k_pe RoPE without local slots", |
| 391 | ) |
| 392 | assert not mla_cache.any(), "non-owner rank wrote the MLA KV cache" |
| 393 | assert not idx_cache.any(), "non-owner rank wrote the indexer KV cache" |
| 394 | assert (topk == -1).all(), "topk buffer not cleared on non-owner rank" |
| 395 | |
| 396 | |
| 397 | @pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS) |
| 398 | @pytest.mark.parametrize("num_tokens", [1, 17, 512]) |
| 399 | def test_fused_norm_rope_no_indexer(num_tokens: int, cfg: ModelConfig): |
| 400 | """Shared (no-indexer) layer: q + kv/MLA only; top-k buffer untouched.""" |
| 401 | torch.manual_seed(1) |
| 402 | dev = "cuda" |
| 403 | max_pos = 8192 |
| 404 | pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) |
| 405 | |
| 406 | q_c = torch.randn(num_tokens, cfg.q_lora, device=dev, dtype=torch.bfloat16) |
| 407 | kv_c = torch.randn(num_tokens, cfg.kv_lora, device=dev, dtype=torch.bfloat16) |
| 408 | k_pe = torch.randn(num_tokens, cfg.rope_dim, device=dev, dtype=torch.bfloat16) |
| 409 | qw = torch.randn(cfg.q_lora, device=dev, dtype=torch.bfloat16) |
| 410 | kvw = torch.randn(cfg.kv_lora, device=dev, dtype=torch.bfloat16) |
| 411 | mla_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev) |
| 412 | |
| 413 | bs = max_pos |
| 414 | mla_cache = torch.zeros( |
| 415 | 1, bs, cfg.kv_lora + cfg.rope_dim, device=dev, dtype=torch.bfloat16 |
| 416 | ) |
| 417 | slot = torch.arange(num_tokens, device=dev, dtype=torch.int64) |
| 418 | topk = torch.full((num_tokens, 2048), 7, device=dev, dtype=torch.int32) |
| 419 | |
| 420 | q_out = K.fused_norm_rope( |
| 421 | pos, |
| 422 | q_c, |
| 423 | qw, |
| 424 | EPS, |
| 425 | kv_c, |
| 426 | kvw, |
| 427 | EPS, |
| 428 | k_pe, |
| 429 | mla_cos_sin, |
| 430 | None, |
| 431 | None, |
| 432 | None, |
| 433 | EPS, |
| 434 | None, |
| 435 | topk, |
| 436 | slot_mapping=slot, |
| 437 | indexer_k_cache=None, |
| 438 | mla_kv_cache=mla_cache, |
| 439 | mla_kv_cache_dtype="auto", |
| 440 | mla_k_scale=None, |
| 441 | has_indexer=False, |
| 442 | index_rope_interleave=False, |
| 443 | ) |
| 444 | |
| 445 | assert_bf16(q_out, rms_norm(q_c, qw), "q_c rmsnorm (no-indexer)") |
| 446 | cache = mla_cache[0, :num_tokens] |
| 447 | assert_bf16(cache[:, : cfg.kv_lora], rms_norm(kv_c, kvw), "MLA kv (no-indexer)") |
| 448 | assert_bf16( |
| 449 | cache[:, cfg.kv_lora :], |
| 450 | rope(k_pe.float(), pos, mla_cos_sin, interleave=True), |
| 451 | "MLA k_pe (no-indexer)", |
| 452 | ) |
| 453 | # Shared layers reuse the previous indexer's top-k: buffer must be untouched. |
| 454 | assert (topk == 7).all(), "topk buffer should be untouched on shared layer" |
| 455 | |
| 456 | |
| 457 | @pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS) |
| 458 | @pytest.mark.parametrize("has_indexer", [False, True]) |
| 459 | def test_fused_norm_rope_materializes_pcp_cache_inputs( |
| 460 | has_indexer: bool, cfg: ModelConfig |
| 461 | ): |
| 462 | """PCP gets local normalized/rotated K rows without direct cache writes.""" |
| 463 | torch.manual_seed(6) |
| 464 | dev = "cuda" |
| 465 | num_tokens = 17 |
| 466 | max_pos = 8192 |
| 467 | pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) |
| 468 | q_c = torch.randn(num_tokens, cfg.q_lora, device=dev, dtype=torch.bfloat16) |
| 469 | kv_c = torch.randn(num_tokens, cfg.kv_lora, device=dev, dtype=torch.bfloat16) |
| 470 | k_pe = torch.randn(num_tokens, cfg.rope_dim, device=dev, dtype=torch.bfloat16) |
| 471 | qw = torch.randn(cfg.q_lora, device=dev, dtype=torch.bfloat16) |
| 472 | kvw = torch.randn(cfg.kv_lora, device=dev, dtype=torch.bfloat16) |
| 473 | ik = ( |
| 474 | torch.randn(num_tokens, cfg.index_head_dim, device=dev, dtype=torch.bfloat16) |
| 475 | if has_indexer |
| 476 | else None |
| 477 | ) |
| 478 | ikw = ( |
| 479 | torch.randn(cfg.index_head_dim, device=dev, dtype=torch.float32) |
| 480 | if has_indexer |
| 481 | else None |
| 482 | ) |
| 483 | ikb = ( |
| 484 | torch.randn(cfg.index_head_dim, device=dev, dtype=torch.float32) |
| 485 | if has_indexer |
| 486 | else None |
| 487 | ) |
| 488 | mla_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev) |
| 489 | idx_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev) if has_indexer else None |
| 490 | q_out = torch.empty_like(q_c) |
| 491 | kv_out = torch.empty_like(kv_c) |
| 492 | kpe_out = torch.empty_like(k_pe) |
| 493 | ik_out = torch.empty_like(ik) if ik is not None else None |
| 494 | topk = torch.full((num_tokens, 2048), 7, device=dev, dtype=torch.int32) |
| 495 | |
| 496 | actual_q = K.fused_norm_rope( |
| 497 | pos, |
| 498 | q_c, |
| 499 | qw, |
| 500 | EPS, |
| 501 | kv_c, |
| 502 | kvw, |
| 503 | EPS, |
| 504 | k_pe, |
| 505 | mla_cos_sin, |
| 506 | ik, |
| 507 | ikw, |
| 508 | ikb, |
| 509 | EPS, |
| 510 | idx_cos_sin, |
| 511 | topk, |
| 512 | has_indexer=has_indexer, |
| 513 | index_rope_interleave=True, |
| 514 | q_c_out=q_out, |
| 515 | kv_c_out=kv_out, |
| 516 | k_pe_out=kpe_out, |
| 517 | index_k_out=ik_out, |
| 518 | ) |
| 519 | |
| 520 | assert actual_q.data_ptr() == q_out.data_ptr() |
| 521 | assert_bf16(actual_q, rms_norm(q_c, qw), "PCP q norm") |
| 522 | assert_bf16(kv_out, rms_norm(kv_c, kvw), "PCP kv norm") |
| 523 | assert_bf16( |
| 524 | kpe_out, |
| 525 | rope(k_pe.float(), pos, mla_cos_sin, interleave=True), |
| 526 | "PCP k_pe RoPE", |
| 527 | ) |
| 528 | if has_indexer: |
| 529 | assert ik is not None and ikw is not None and ikb is not None |
| 530 | assert ik_out is not None and idx_cos_sin is not None |
| 531 | ik_ref = rope(layer_norm(ik, ikw, ikb), pos, idx_cos_sin, interleave=True) |
| 532 | assert_bf16(ik_out, ik_ref, "PCP indexer-K") |
| 533 | |
| 534 | |
| 535 | @pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS) |
| 536 | @pytest.mark.parametrize("num_tokens", [1, 4, 17, 512]) |
| 537 | def test_fused_norm_rope_ds_mla(num_tokens: int, cfg: ModelConfig): |
| 538 | """fp8_ds_mla MLA cache layout (FlashMLA sparse, bf16-query path; SM90/SM100). |
| 539 | |
| 540 | Per-token 656-byte entry: 512 fp8 NoPE (4 per-128 tiles, power-of-two scales |
| 541 | stored as float32) | 4 float32 scales | 64 bf16 (unquantized) RoPE. |
| 542 | """ |
| 543 | torch.manual_seed(5) |
| 544 | dev = "cuda" |
| 545 | max_pos = 8192 |
| 546 | pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos |
| 547 | |
| 548 | q_c = torch.randn(num_tokens, cfg.q_lora, device=dev, dtype=torch.bfloat16) |
| 549 | kv_c = torch.randn(num_tokens, cfg.kv_lora, device=dev, dtype=torch.bfloat16) |
| 550 | k_pe = torch.randn(num_tokens, cfg.rope_dim, device=dev, dtype=torch.bfloat16) |
| 551 | qw = torch.randn(cfg.q_lora, device=dev, dtype=torch.bfloat16) |
| 552 | kvw = torch.randn(cfg.kv_lora, device=dev, dtype=torch.bfloat16) |
| 553 | mla_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev) |
| 554 | |
| 555 | bs = max_pos |
| 556 | mla_cache = torch.zeros(1, bs, 656, device=dev, dtype=torch.uint8) |
| 557 | slot = torch.arange(num_tokens, device=dev, dtype=torch.int64) |
| 558 | topk = torch.full((num_tokens, 2048), 7, device=dev, dtype=torch.int32) |
| 559 | |
| 560 | q_out = K.fused_norm_rope( |
| 561 | pos, |
| 562 | q_c, |
| 563 | qw, |
| 564 | EPS, |
| 565 | kv_c, |
| 566 | kvw, |
| 567 | EPS, |
| 568 | k_pe, |
| 569 | mla_cos_sin, |
| 570 | None, |
| 571 | None, |
| 572 | None, |
| 573 | EPS, |
| 574 | None, |
| 575 | topk, |
| 576 | slot_mapping=slot, |
| 577 | indexer_k_cache=None, |
| 578 | mla_kv_cache=mla_cache, |
| 579 | mla_kv_cache_dtype="fp8_ds_mla", |
| 580 | mla_k_scale=None, |
| 581 | has_indexer=False, |
| 582 | index_rope_interleave=False, |
| 583 | ) |
| 584 | |
| 585 | assert_bf16(q_out, rms_norm(q_c, qw), "q_c rmsnorm (ds_mla)") |
| 586 | |
| 587 | kv_ref = rms_norm(kv_c, kvw) # [N, 512] fp32 |
| 588 | kpe_ref = rope(k_pe.float(), pos, mla_cos_sin, interleave=True) # [N, 64] |
| 589 | tiles = kv_ref.view(num_tokens, 4, 128) |
| 590 | ref_scale = torch.clamp(tiles.abs().amax(dim=-1) / FP8_MAX, min=1e-4) |
| 591 | ref_scale = torch.exp2(torch.ceil(torch.log2(ref_scale))) |
| 592 | ref_nope = (tiles / ref_scale[..., None]).reshape(num_tokens, cfg.kv_lora).to(FP8) |
| 593 | |
| 594 | cache = mla_cache[0, :num_tokens] # [N, 656] uint8 |
| 595 | nope = cache[:, : cfg.kv_lora].view(FP8) |
| 596 | scales = cache.view(torch.float32)[:, cfg.kv_lora // 4 : cfg.kv_lora // 4 + 4] |
| 597 | rope_off = cfg.kv_lora // 2 + 8 |
| 598 | rope_vals = cache.view(torch.bfloat16)[:, rope_off : rope_off + cfg.rope_dim] |
| 599 | |
| 600 | torch.testing.assert_close(scales, ref_scale, rtol=0, atol=0) |
| 601 | assert_fp8(nope, ref_nope, "ds_mla NoPE fp8") |
| 602 | assert_bf16(rope_vals, kpe_ref, "ds_mla RoPE bf16") |
| 603 | # No indexer on this call: top-k buffer must be untouched. |
| 604 | assert (topk == 7).all(), "topk buffer should be untouched (no indexer)" |
| 605 | |
| 606 | |
| 607 | E2M1_MAGNITUDES = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0] |
| 608 | |
| 609 | |
| 610 | def quantize_to_e2m1(x: torch.Tensor) -> torch.Tensor: |
| 611 | """Round to nearest e2m1, saturating to +-6 like cvt.rn.satfinite.e2m1x2.f32.""" |
| 612 | mags = torch.tensor(E2M1_MAGNITUDES, dtype=torch.float32, device=x.device) |
| 613 | a = x.float().abs().clamp_max(6.0) |
| 614 | mids = (mags[:-1] + mags[1:]) / 2 |
| 615 | code = torch.bucketize(a, mids, right=True) |
| 616 | on_tie = (a.unsqueeze(-1) == mids).any(dim=-1) |
| 617 | tie_code = torch.bucketize(a, mids, right=False) |
| 618 | code = torch.where(on_tie, tie_code + (tie_code & 1), code) |
| 619 | return (torch.signbit(x).to(torch.uint8) << 3) | code.to(torch.uint8) |
| 620 | |
| 621 | |
| 622 | def nvfp4_sf_byte(s: torch.Tensor) -> torch.Tensor: |
| 623 | """Scale-factor byte permutation shared with FlashMLA: 8*(s&3) + (s>>2).""" |
| 624 | return 8 * (s % 4) + (s // 4) |
| 625 | |
| 626 | |
| 627 | @pytest.mark.skipif( |
| 628 | not current_platform.is_device_capability(100), |
| 629 | reason="nvfp4_ds_mla requires SM100 (Blackwell)", |
| 630 | ) |
| 631 | @pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS) |
| 632 | @pytest.mark.parametrize("num_tokens", [1, 4, 17, 512]) |
| 633 | def test_fused_norm_rope_nvfp4_ds_mla(num_tokens: int, cfg: ModelConfig): |
| 634 | """nvfp4_ds_mla MLA cache layout (FlashMLA sparse, SM100 only). |
| 635 | |
| 636 | Per-token 352-byte entry: 256 B of 512 e2m1 NoPE packed 2/byte (low nibble |
| 637 | = even element) | 64 B unscaled e4m3 RoPE | 32 B byte-permuted e4m3 tile |
| 638 | scales, one per 16 NoPE elements. |
| 639 | """ |
| 640 | torch.manual_seed(5) |
| 641 | dev = "cuda" |
| 642 | max_pos = 8192 |
| 643 | pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos |
| 644 | |
| 645 | q_c = torch.randn(num_tokens, cfg.q_lora, device=dev, dtype=torch.bfloat16) |
| 646 | kv_c = torch.randn(num_tokens, cfg.kv_lora, device=dev, dtype=torch.bfloat16) |
| 647 | # Tile 0 is deliberately low-magnitude: amax/6 lands between two e4m3 |
| 648 | # subnormals (spaced a flat 2^-9 there), so round-to-nearest picks a scale |
| 649 | # BELOW amax/6 and the tile's peak saturates at +-6. Asserted below, so a |
| 650 | # future change to the scale rule cannot silently stop covering this. |
| 651 | # amax must land in (6, 9) * 2^-9 after rms_norm; 7.5 sits mid-band, with |
| 652 | # margin for the per-token RMS to vary. |
| 653 | kv_c[:, :16] = ( |
| 654 | 7.5 * (2.0**-9) * torch.linspace(0.2, 1.0, 16, device=dev, dtype=torch.bfloat16) |
| 655 | ) |
| 656 | k_pe = torch.randn(num_tokens, cfg.rope_dim, device=dev, dtype=torch.bfloat16) |
| 657 | qw = torch.randn(cfg.q_lora, device=dev, dtype=torch.bfloat16) |
| 658 | kvw = torch.ones(cfg.kv_lora, device=dev, dtype=torch.bfloat16) |
| 659 | mla_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev) |
| 660 | |
| 661 | bs = max_pos |
| 662 | mla_cache = torch.zeros(1, bs, 352, device=dev, dtype=torch.uint8) |
| 663 | slot = torch.arange(num_tokens, device=dev, dtype=torch.int64) |
| 664 | topk = torch.full((num_tokens, 2048), 7, device=dev, dtype=torch.int32) |
| 665 | |
| 666 | q_out = K.fused_norm_rope( |
| 667 | pos, |
| 668 | q_c, |
| 669 | qw, |
| 670 | EPS, |
| 671 | kv_c, |
| 672 | kvw, |
| 673 | EPS, |
| 674 | k_pe, |
| 675 | mla_cos_sin, |
| 676 | None, |
| 677 | None, |
| 678 | None, |
| 679 | EPS, |
| 680 | None, |
| 681 | topk, |
| 682 | slot_mapping=slot, |
| 683 | indexer_k_cache=None, |
| 684 | mla_kv_cache=mla_cache, |
| 685 | mla_kv_cache_dtype="nvfp4_ds_mla", |
| 686 | mla_k_scale=None, |
| 687 | has_indexer=False, |
| 688 | index_rope_interleave=False, |
| 689 | ) |
| 690 | |
| 691 | assert_bf16(q_out, rms_norm(q_c, qw), "q_c rmsnorm (nvfp4_ds_mla)") |
| 692 | |
| 693 | kv_ref = rms_norm(kv_c, kvw) # [N, 512] fp32 |
| 694 | kpe_ref = rope(k_pe.float(), pos, mla_cos_sin, interleave=True) # [N, 64] |
| 695 | tiles = kv_ref.view(num_tokens, 32, 16) |
| 696 | amax = tiles.abs().amax(dim=-1) |
| 697 | scale_target = torch.clamp_min(amax / 6.0, 2.0**-9) |
| 698 | ref_scale = scale_target.to(FP8) # round-to-nearest |
| 699 | |
| 700 | # The crafted tile must actually exercise the saturating path. |
| 701 | assert (ref_scale[:, 0].float() < scale_target[:, 0]).all(), ( |
| 702 | "tile 0 should round its scale DOWN; adjust the crafted magnitude" |
| 703 | ) |
| 704 | assert (amax[:, 0] / ref_scale[:, 0].float() > 6.0).all(), ( |
| 705 | "tile 0 should saturate e2m1; adjust the crafted magnitude" |
| 706 | ) |
| 707 | |
| 708 | ref_codes = quantize_to_e2m1(tiles / ref_scale.float().unsqueeze(-1)) |
| 709 | ref_codes = ref_codes.reshape(num_tokens, cfg.kv_lora) |
| 710 | |
| 711 | cache = mla_cache[0, :num_tokens] # [N, 352] uint8 |
| 712 | packed = cache[:, : cfg.kv_lora // 2] |
| 713 | got_codes = torch.stack([packed & 0xF, packed >> 4], dim=-1) |
| 714 | got_codes = got_codes.reshape(num_tokens, cfg.kv_lora) # low nibble = even elem |
| 715 | torch.testing.assert_close(got_codes, ref_codes, rtol=0, atol=0) |
| 716 | |
| 717 | got_rope = cache[:, cfg.kv_lora // 2 : cfg.kv_lora // 2 + cfg.rope_dim].view(FP8) |
| 718 | assert_fp8(got_rope, kpe_ref.to(FP8), "nvfp4_ds_mla RoPE e4m3") |
| 719 | |
| 720 | perm = nvfp4_sf_byte(torch.arange(32, device=dev)) |
| 721 | got_scale = cache[:, cfg.kv_lora // 2 + cfg.rope_dim :].view(FP8)[:, perm] |
| 722 | torch.testing.assert_close(got_scale.float(), ref_scale.float(), rtol=0, atol=0) |
| 723 | |
| 724 | # No indexer on this call: top-k buffer must be untouched. |
| 725 | assert (topk == 7).all(), "topk buffer should be untouched (no indexer)" |
| 726 | |
| 727 | |
| 728 | def test_fused_norm_rope_supports_large_token_count(): |
| 729 | """Keep the token count off CUDA grid-y at its 65,536-block boundary.""" |
| 730 | num_tokens = 65536 |
| 731 | dev = "cuda" |
| 732 | dtype = torch.bfloat16 |
| 733 | positions = torch.zeros(num_tokens, device=dev, dtype=torch.int64) |
| 734 | q_c = torch.ones((num_tokens, 1), device=dev, dtype=dtype) |
| 735 | kv_c = torch.ones((num_tokens, 1), device=dev, dtype=dtype) |
| 736 | k_pe = torch.ones((num_tokens, 2), device=dev, dtype=dtype) |
| 737 | norm_w = torch.ones(1, device=dev, dtype=dtype) |
| 738 | cos_sin = torch.tensor([[1.0, 0.0]], device=dev, dtype=torch.float32) |
| 739 | topk = torch.empty((num_tokens, 1), device=dev, dtype=torch.int32) |
| 740 | slot_mapping = torch.arange(num_tokens, device=dev, dtype=torch.int64) |
| 741 | mla_cache = torch.empty((1, num_tokens, 3), device=dev, dtype=dtype) |
| 742 | |
| 743 | q_out = K.fused_norm_rope( |
| 744 | positions, |
| 745 | q_c, |
| 746 | norm_w, |
| 747 | EPS, |
| 748 | kv_c, |
| 749 | norm_w, |
| 750 | EPS, |
| 751 | k_pe, |
| 752 | cos_sin, |
| 753 | None, |
| 754 | None, |
| 755 | None, |
| 756 | EPS, |
| 757 | None, |
| 758 | topk, |
| 759 | slot_mapping=slot_mapping, |
| 760 | mla_kv_cache=mla_cache, |
| 761 | has_indexer=False, |
| 762 | ) |
| 763 | |
| 764 | rows = torch.tensor([0, num_tokens - 1], device=dev) |
| 765 | assert_bf16(q_out[rows], rms_norm(q_c[rows], norm_w), "large-token q norm") |
| 766 | |
| 767 | |
| 768 | # ── fused_q ────────────────────────────────────────────────────────────────── |
| 769 | |
| 770 | |
| 771 | @pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS) |
| 772 | @pytest.mark.parametrize("num_tokens", [1, 4, 17, 512, 4096]) |
| 773 | @pytest.mark.parametrize("index_interleave", [True, False]) |
| 774 | def test_fused_q(num_tokens: int, index_interleave: bool, cfg: ModelConfig): |
| 775 | torch.manual_seed(2) |
| 776 | dev = "cuda" |
| 777 | max_pos = 8192 |
| 778 | pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos |
| 779 | |
| 780 | q_pe = torch.randn( |
| 781 | num_tokens, cfg.num_heads, cfg.rope_dim, device=dev, dtype=torch.bfloat16 |
| 782 | ) |
| 783 | ql_nope = torch.randn( |
| 784 | num_tokens, cfg.num_heads, cfg.kv_lora, device=dev, dtype=torch.bfloat16 |
| 785 | ) |
| 786 | index_q = torch.randn( |
| 787 | num_tokens, |
| 788 | cfg.index_heads, |
| 789 | cfg.index_head_dim, |
| 790 | device=dev, |
| 791 | dtype=torch.bfloat16, |
| 792 | ) |
| 793 | index_w = torch.randn(num_tokens, cfg.index_heads, device=dev, dtype=torch.float32) |
| 794 | q_scale = torch.tensor([0.37], device=dev, dtype=torch.float32) |
| 795 | softmax_scale = cfg.index_head_dim**-0.5 |
| 796 | head_scale = cfg.index_heads**-0.5 |
| 797 | q_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev) # q_pe: interleaved |
| 798 | idx_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev) |
| 799 | |
| 800 | iq_fp8, iw_out, mqa = K.fused_q( |
| 801 | pos, |
| 802 | q_pe, |
| 803 | q_cos_sin, |
| 804 | index_q, |
| 805 | idx_cos_sin, |
| 806 | ql_nope, |
| 807 | q_scale, |
| 808 | index_w, |
| 809 | softmax_scale, |
| 810 | head_scale, |
| 811 | has_indexer=True, |
| 812 | index_rope_interleave=index_interleave, |
| 813 | ) |
| 814 | |
| 815 | s = q_scale.item() |
| 816 | # MQA query: [ql_nope | q_pe RoPE'd (interleaved)], per-tensor fp8. |
| 817 | mqa_nope_ref = (ql_nope.float() / s).to(FP8) |
| 818 | qpe_ref = rope( |
| 819 | q_pe.float(), |
| 820 | pos.unsqueeze(-1).expand(num_tokens, cfg.num_heads), |
| 821 | q_cos_sin, |
| 822 | interleave=True, |
| 823 | ) |
| 824 | mqa_pe_ref = (qpe_ref / s).to(FP8) |
| 825 | assert_fp8(mqa[:, :, : cfg.kv_lora], mqa_nope_ref, "mqa ql_nope") |
| 826 | assert_fp8(mqa[:, :, cfg.kv_lora :], mqa_pe_ref, "mqa q_pe") |
| 827 | |
| 828 | # Indexer-Q: RoPE + UE8M0 fp8 quant; index weights fold in q-scale. |
| 829 | iq_ref = rope( |
| 830 | index_q.float(), |
| 831 | pos.unsqueeze(-1).expand(num_tokens, cfg.index_heads), |
| 832 | idx_cos_sin, |
| 833 | interleave=index_interleave, |
| 834 | ) |
| 835 | q_ref, scale_ref = ue8m0_quant(iq_ref) |
| 836 | assert_fp8(iq_fp8, q_ref, "indexer-Q fp8") |
| 837 | iw_ref = index_w * scale_ref * softmax_scale * head_scale |
| 838 | torch.testing.assert_close(iw_out, iw_ref, rtol=1e-3, atol=1e-3) |
| 839 | |
| 840 | |
| 841 | @pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS) |
| 842 | @pytest.mark.parametrize("num_tokens", [1, 17, 512]) |
| 843 | def test_fused_q_no_indexer(num_tokens: int, cfg: ModelConfig): |
| 844 | torch.manual_seed(3) |
| 845 | dev = "cuda" |
| 846 | max_pos = 8192 |
| 847 | pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) |
| 848 | q_pe = torch.randn( |
| 849 | num_tokens, cfg.num_heads, cfg.rope_dim, device=dev, dtype=torch.bfloat16 |
| 850 | ) |
| 851 | ql_nope = torch.randn( |
| 852 | num_tokens, cfg.num_heads, cfg.kv_lora, device=dev, dtype=torch.bfloat16 |
| 853 | ) |
| 854 | q_scale = torch.tensor([0.5], device=dev, dtype=torch.float32) |
| 855 | q_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev) |
| 856 | _, _, mqa = K.fused_q( |
| 857 | pos, |
| 858 | q_pe, |
| 859 | q_cos_sin, |
| 860 | None, |
| 861 | None, |
| 862 | ql_nope, |
| 863 | q_scale, |
| 864 | None, |
| 865 | 0.0, |
| 866 | 0.0, |
| 867 | has_indexer=False, |
| 868 | index_rope_interleave=False, |
| 869 | ) |
| 870 | s = q_scale.item() |
| 871 | assert_fp8(mqa[:, :, : cfg.kv_lora], (ql_nope.float() / s).to(FP8), "mqa ql_nope") |
| 872 | qpe_ref = rope( |
| 873 | q_pe.float(), |
| 874 | pos.unsqueeze(-1).expand(num_tokens, cfg.num_heads), |
| 875 | q_cos_sin, |
| 876 | interleave=True, |
| 877 | ) |
| 878 | assert_fp8(mqa[:, :, cfg.kv_lora :], (qpe_ref / s).to(FP8), "mqa q_pe") |
| 879 | |
| 880 | |
| 881 | @pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS) |
| 882 | @pytest.mark.parametrize("num_tokens", [1, 17, 512]) |
| 883 | @pytest.mark.parametrize("has_indexer", [True, False]) |
| 884 | @pytest.mark.parametrize("use_mha", [True, False]) |
| 885 | @pytest.mark.parametrize( |
| 886 | "index_interleave,cache_dtype", [(False, torch.float32), (True, torch.bfloat16)] |
| 887 | ) |
| 888 | @pytest.mark.parametrize("capture", [False, True]) |
| 889 | def test_fused_q_bf16_query( |
| 890 | num_tokens: int, |
| 891 | has_indexer: bool, |
| 892 | use_mha: bool, |
| 893 | index_interleave: bool, |
| 894 | cache_dtype: torch.dtype, |
| 895 | capture: bool, |
| 896 | cfg: ModelConfig, |
| 897 | ): |
| 898 | """bf16-query path (FlashMLA sparse, SM90/SM100): only the RoPE'd q_pe is |
| 899 | produced; MHA does not need an absorbed ql_nope projection.""" |
| 900 | torch.manual_seed(6) |
| 901 | dev = "cuda" |
| 902 | max_pos = 8192 |
| 903 | pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) % max_pos |
| 904 | |
| 905 | # Attention takes q_pe as a strided slice of the full Q projection. |
| 906 | q_pe = torch.randn( |
| 907 | num_tokens, cfg.num_heads, 4 * cfg.rope_dim, device=dev, dtype=torch.bfloat16 |
| 908 | )[..., -cfg.rope_dim :] |
| 909 | ql_nope = torch.randn( |
| 910 | num_tokens, cfg.num_heads, cfg.kv_lora, device=dev, dtype=torch.bfloat16 |
| 911 | ) |
| 912 | q_scale = torch.tensor([0.37], device=dev, dtype=torch.float32) |
| 913 | q_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev).to(cache_dtype) |
| 914 | |
| 915 | index_q = index_w = idx_cos_sin = None |
| 916 | if has_indexer: |
| 917 | index_q = torch.randn( |
| 918 | num_tokens, |
| 919 | cfg.index_heads, |
| 920 | cfg.index_head_dim, |
| 921 | device=dev, |
| 922 | dtype=torch.bfloat16, |
| 923 | ) |
| 924 | index_w = torch.randn( |
| 925 | num_tokens, cfg.index_heads, device=dev, dtype=torch.float32 |
| 926 | ) |
| 927 | idx_cos_sin = make_cos_sin(max_pos, cfg.rope_dim, dev).to(cache_dtype) |
| 928 | |
| 929 | run = partial( |
| 930 | K.fused_q, |
| 931 | pos, |
| 932 | q_pe, |
| 933 | q_cos_sin, |
| 934 | index_q, |
| 935 | idx_cos_sin, |
| 936 | q_pe if use_mha else ql_nope, |
| 937 | q_scale, |
| 938 | index_w, |
| 939 | cfg.index_head_dim**-0.5, |
| 940 | cfg.index_heads**-0.5, |
| 941 | has_indexer=has_indexer, |
| 942 | index_rope_interleave=index_interleave, |
| 943 | quantize_mqa=False, |
| 944 | ) |
| 945 | |
| 946 | iq_fp8, iw_out, q_pe_out = run() |
| 947 | if capture: |
| 948 | graph = torch.cuda.CUDAGraph() |
| 949 | with torch.cuda.graph(graph): |
| 950 | iq_fp8, iw_out, q_pe_out = run() |
| 951 | q_pe.mul_(0.5) |
| 952 | if index_q is not None: |
| 953 | index_q.mul_(1.5) |
| 954 | graph.replay() |
| 955 | |
| 956 | # MQA query: only the RoPE'd q_pe, bf16, unquantized. |
| 957 | assert q_pe_out.dtype == torch.bfloat16 |
| 958 | assert q_pe_out.shape == (num_tokens, cfg.num_heads, cfg.rope_dim) |
| 959 | qpe_ref = rope( |
| 960 | q_pe.float(), |
| 961 | pos.unsqueeze(-1).expand(num_tokens, cfg.num_heads), |
| 962 | q_cos_sin, |
| 963 | interleave=True, |
| 964 | ) |
| 965 | assert_bf16(q_pe_out, qpe_ref, "bf16 q_pe RoPE") |
| 966 | |
| 967 | # Indexer-Q is unchanged on this path (still UE8M0 fp8 + folded weights). |
| 968 | if has_indexer: |
| 969 | assert index_q is not None |
| 970 | iq_ref = rope( |
| 971 | index_q.float(), |
| 972 | pos.unsqueeze(-1).expand(num_tokens, cfg.index_heads), |
| 973 | idx_cos_sin, |
| 974 | interleave=index_interleave, |
| 975 | ) |
| 976 | q_ref, scale_ref = ue8m0_quant(iq_ref) |
| 977 | assert_fp8(iq_fp8, q_ref, "indexer-Q fp8 (bf16-query path)") |
| 978 | iw_ref = ( |
| 979 | index_w * scale_ref * (cfg.index_head_dim**-0.5) * (cfg.index_heads**-0.5) |
| 980 | ) |
| 981 | torch.testing.assert_close(iw_out, iw_ref, rtol=1e-3, atol=1e-3) |
| 982 | |
| 983 | |
| 984 | def test_fused_q_triton_supports_large_token_count(): |
| 985 | """Keep the token count off CUDA grid-y in the Triton fallback. |
| 986 | |
| 987 | The minimal dimensions also bypass CuTeDSL on SM100. |
| 988 | """ |
| 989 | num_tokens = 65536 |
| 990 | dev = "cuda" |
| 991 | dtype = torch.bfloat16 |
| 992 | positions = torch.zeros(num_tokens, device=dev, dtype=torch.int64) |
| 993 | q_pe = torch.ones((num_tokens, 1, 2), device=dev, dtype=dtype) |
| 994 | ql_nope = torch.ones((num_tokens, 1, 1), device=dev, dtype=dtype) |
| 995 | cos_sin = torch.tensor([[1.0, 0.0]], device=dev, dtype=torch.float32) |
| 996 | q_scale = torch.ones(1, device=dev, dtype=torch.float32) |
| 997 | |
| 998 | _, _, mqa_q = K.fused_q( |
| 999 | positions, |
| 1000 | q_pe, |
| 1001 | cos_sin, |
| 1002 | None, |
| 1003 | None, |
| 1004 | ql_nope, |
| 1005 | q_scale, |
| 1006 | None, |
| 1007 | 0.0, |
| 1008 | 0.0, |
| 1009 | has_indexer=False, |
| 1010 | ) |
| 1011 | |
| 1012 | rows = torch.tensor([0, num_tokens - 1], device=dev) |
| 1013 | ref = torch.cat([ql_nope, q_pe], dim=-1).to(FP8) |
| 1014 | assert_fp8(mqa_q[rows], ref[rows], "large-token fused Q") |
| 1015 | |
| 1016 | |
| 1017 | # ── fused_eh_norm (MTP) ────────────────────────────────────────────────────── |
| 1018 | |
| 1019 | |
| 1020 | @pytest.mark.parametrize("cfg", MODEL_CONFIGS, ids=MODEL_IDS) |
| 1021 | @pytest.mark.parametrize("num_tokens", [1, 4, 17, 512]) |
| 1022 | def test_fused_eh_norm(num_tokens: int, cfg: ModelConfig): |
| 1023 | torch.manual_seed(4) |
| 1024 | dev = "cuda" |
| 1025 | # Mix in a position-0 token to exercise the embeds-zeroing branch. |
| 1026 | pos = torch.arange(num_tokens, device=dev, dtype=torch.int64) |
| 1027 | pos[0] = 0 |
| 1028 | embeds = torch.randn(num_tokens, cfg.hidden, device=dev, dtype=torch.bfloat16) |
| 1029 | prev = torch.randn(num_tokens, cfg.hidden, device=dev, dtype=torch.bfloat16) |
| 1030 | ew = torch.randn(cfg.hidden, device=dev, dtype=torch.bfloat16) |
| 1031 | hw = torch.randn(cfg.hidden, device=dev, dtype=torch.bfloat16) |
| 1032 | |
| 1033 | out = K.fused_eh_norm(pos, embeds, prev, ew, hw, EPS) |
| 1034 | |
| 1035 | masked = torch.where(pos.unsqueeze(-1) == 0, torch.zeros_like(embeds), embeds) |
| 1036 | ref = torch.cat([rms_norm(masked, ew), rms_norm(prev, hw)], dim=-1) |
| 1037 | assert out.shape == (num_tokens, 2 * cfg.hidden) |
| 1038 | assert_bf16(out, ref, "eh_norm") |