A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 11:15 UTC
[DSv4] Use `cvt` PTX for FP32->FP4 conversion (#41015)
Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>
Thien Tran committed
Apr 30, 2026 at 06:16 UTC
296741d0257107a9d0301409005c85d38bb247bc
4 files changed
+344
-62
tests/kernels/test_compressor_kv_cache.py
+228
-4
@@ -3,12 +3,11 @@
3
"""
4
Round-trip tests for compressor → FP8 quant + KV cache insert → gather + dequant.
5
6
-Two paths tested:
6
+Four test functions cover five paths:
7
A) DeepseekV4 Attention: head_dim=512 (448 FP8 nope + 64 bf16 rope), quant_block=64
8
B) Indexer: head_dim=128 (all FP8), quant_block=128
9
-
10
-These serve as golden references for validating the future fused
11
-compressor+quant+cache kernel.
9
+ C) DeepseekV4 Attention magnitude range: correctness across small/large values
10
+ D) Indexer fused Triton kernel: compress+norm+rope+quant+insert
11
"""
12
13
import math
@@ -21,6 +20,12 @@ from vllm.v1.attention.ops.deepseek_v4_ops import (
20
dequantize_and_gather_k_cache,
21
quantize_and_insert_k_cache,
22
)
23
+from vllm.v1.attention.ops.deepseek_v4_ops.fused_compress_quant_cache import (
24
+ _fused_kv_compress_norm_rope_insert_indexer_attn,
25
+ _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn,
26
+)
27
+
28
+from .test_fused_indexer_q_rope_quant import quantize_to_mxfp4
29
30
31
def _ue8m0_reference(x: torch.Tensor, block_size: int, fp8_max: float):
@@ -309,3 +314,222 @@ def test_deepseek_v4_quant_magnitude_range():
314
f"Token {t}: rel_err={rel_err:.4f}, abs_diff={abs_diff:.6f}, "
315
f"magnitude={magnitude:.4f}"
316
)
317
+
318
+
319
+# ── Test D: Indexer fused K-cache insert (Triton kernels) ────────────────────
320
+#
321
+# Both kernels share the same Triton signature; use_fp4 selects between them.
322
+# Full pipeline: state-cache gather → softmax-weighted compress → RMSNorm →
323
+# GPT-J RoPE → quant (MXFP4 or FP8) → paged cache insert.
324
+
325
+
326
+def _reference_kv_compress_norm_rope(
327
+ state_cache: torch.Tensor,
328
+ block_table: torch.Tensor,
329
+ positions: torch.Tensor,
330
+ rms_weight: torch.Tensor,
331
+ cos_sin_cache: torch.Tensor,
332
+ compress_ratio: int = 1,
333
+ overlap: int = 0,
334
+ use_fp4: bool = False,
335
+ rms_eps: float = 1e-6,
336
+ fp8_max: float = 448.0,
337
+):
338
+ """Compress → RMSNorm → GPT-J RoPE → quantize.
339
+
340
+ Gathers (1+overlap)*compress_ratio state entries per output token, applies
341
+ per-element softmax over the scores, and computes the weighted kv sum.
342
+ Returns (quantized_values, scale) matching the kernel's output layout.
343
+ """
344
+ device = state_cache.device
345
+ head_dim = rms_weight.shape[0]
346
+ rope_dim = cos_sin_cache.shape[-1]
347
+ state_block_size = state_cache.shape[1]
348
+ state_width = state_cache.shape[-1] // 2
349
+ nope_dim = head_dim - rope_dim
350
+ total = (1 + overlap) * compress_ratio
351
+ results = []
352
+ for pos in positions.tolist():
353
+ src = torch.arange(pos - total + 1, pos + 1, dtype=torch.int64, device=device)
354
+ valid = src >= 0
355
+ idx = src.clamp(min=0)
356
+ pages = block_table[0, idx // state_block_size]
357
+ offsets = idx % state_block_size
358
+ raw = state_cache[pages, offsets].float() # [total, state_dim]
359
+
360
+ # Group 0 (tokens 0..cr-1): kv[:H], score[SW:SW+H]
361
+ # Group 1 (tokens cr..2cr-1): kv[H:2H], score[SW+H:SW+2H]
362
+ if overlap:
363
+ sw = state_width
364
+ g0_kv = raw[:compress_ratio, :head_dim]
365
+ g1_kv = raw[compress_ratio:, head_dim : 2 * head_dim]
366
+ g0_scores = raw[:compress_ratio, sw : sw + head_dim]
367
+ g1_scores = raw[compress_ratio:, sw + head_dim : sw + 2 * head_dim]
368
+ kv = torch.cat([g0_kv, g1_kv])
369
+ scores = torch.cat([g0_scores, g1_scores])
370
+ else:
371
+ kv = raw[:, :head_dim]
372
+ scores = raw[:, state_width : state_width + head_dim]
373
+
374
+ scores[~valid] = float("-inf")
375
+ kv[~valid] = 0.0
376
+ weights = torch.softmax(scores, dim=0)
377
+ compressed = (kv * weights).sum(dim=0) # [H]
378
+ var = (compressed * compressed).mean()
379
+ normed = compressed * torch.rsqrt(var + rms_eps) * rms_weight.float()
380
+ compressed_pos = (pos // compress_ratio) * compress_ratio
381
+ cos, sin = cos_sin_cache[compressed_pos].float().chunk(2)
382
+ nope, rope = normed.split([nope_dim, rope_dim])
383
+ rope = torch.stack(
384
+ [rope[0::2] * cos - rope[1::2] * sin, rope[1::2] * cos + rope[0::2] * sin],
385
+ dim=-1,
386
+ ).reshape(rope_dim)
387
+ results.append(torch.cat([nope, rope]).to(state_cache.dtype))
388
+ result = torch.stack(results)
389
+
390
+ if use_fp4:
391
+ return quantize_to_mxfp4(result)
392
+ else:
393
+ pairs = [
394
+ _ue8m0_reference(result[t], head_dim, fp8_max) for t in range(len(result))
395
+ ]
396
+ quants, scales = zip(*pairs)
397
+ return torch.stack(quants), torch.cat(scales)
398
+
399
+
400
+@pytest.mark.parametrize("num_tokens", [1, 7, 32])
401
+@pytest.mark.parametrize("kv_block_size", [16, 32])
402
+@pytest.mark.parametrize("use_fp4", [False, True])
403
+def test_fused_kv_insert_indexer(num_tokens: int, kv_block_size: int, use_fp4: bool):
404
+ """Fused K compress+norm+rope+quant+insert for the indexer KV cache."""
405
+ HEAD_DIM = 128
406
+ ROPE_DIM = 64
407
+ BLOCK_SIZE = 16
408
+ RMS_EPS = 1e-6
409
+ FP8_MAX = 448.0
410
+
411
+ device = "cuda"
412
+ torch.manual_seed(42)
413
+ compress_ratio = 4
414
+
415
+ if use_fp4:
416
+ TOKEN_STRIDE = HEAD_DIM // 2 # packed nibbles: 64 bytes
417
+ SCALE_DIM = HEAD_DIM // 32 # ue8m0 bytes: 4
418
+ QUANT_BLOCK = 32
419
+ kernel = _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn
420
+ else:
421
+ TOKEN_STRIDE = HEAD_DIM # FP8 bytes: 128
422
+ SCALE_DIM = 4 # 1 float32: 4 bytes
423
+ QUANT_BLOCK = HEAD_DIM
424
+ kernel = _fused_kv_compress_norm_rope_insert_indexer_attn
425
+
426
+ # overlap=1 whenever compress_ratio==4, matching DeepseekCompressor logic.
427
+ overlap = 1 if compress_ratio == 4 else 0
428
+ coff = 1 + overlap # multiplier for state_dim per entry
429
+
430
+ num_pages = (compress_ratio * num_tokens - 1) // BLOCK_SIZE + 2
431
+ state_cache = torch.randn(
432
+ num_pages,
433
+ BLOCK_SIZE,
434
+ 2 * coff * HEAD_DIM, # kv_state + score_state, each coff*HEAD_DIM wide
435
+ dtype=torch.bfloat16,
436
+ device=device,
437
+ )
438
+ block_table = torch.arange(num_pages, dtype=torch.int32, device=device).unsqueeze(0)
439
+ token_to_req = torch.zeros(num_tokens, dtype=torch.int32, device=device)
440
+ slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device)
441
+ positions = torch.arange(
442
+ compress_ratio - 1,
443
+ compress_ratio * num_tokens,
444
+ compress_ratio,
445
+ dtype=torch.int64,
446
+ device=device,
447
+ )
448
+ rms_weight = torch.randn(HEAD_DIM, dtype=torch.bfloat16, device=device)
449
+ cos_sin_cache = torch.randn(compress_ratio * num_tokens, ROPE_DIM, device=device)
450
+
451
+ kv_n_blocks = (num_tokens + kv_block_size - 1) // kv_block_size + 1
452
+ kv_cache = torch.zeros(
453
+ kv_n_blocks,
454
+ kv_block_size * (TOKEN_STRIDE + SCALE_DIM),
455
+ dtype=torch.uint8,
456
+ device=device,
457
+ )
458
+
459
+ kernel[(num_tokens,)](
460
+ state_cache,
461
+ state_cache.stride(0),
462
+ state_cache.stride(1),
463
+ token_to_req,
464
+ positions,
465
+ slot_mapping,
466
+ block_table,
467
+ block_table.stride(0),
468
+ BLOCK_SIZE,
469
+ rms_weight,
470
+ RMS_EPS,
471
+ cos_sin_cache,
472
+ cos_sin_cache.stride(0),
473
+ kv_cache,
474
+ slot_mapping,
475
+ kv_block_size,
476
+ HEAD_SIZE=HEAD_DIM,
477
+ TRITON_BLOCK_SIZE=HEAD_DIM,
478
+ STATE_WIDTH=coff * HEAD_DIM,
479
+ COMPRESS_RATIO=compress_ratio,
480
+ OVERLAP=overlap,
481
+ ROPE_HEAD_DIM=ROPE_DIM,
482
+ FP8_MAX=FP8_MAX,
483
+ QUANT_BLOCK=QUANT_BLOCK,
484
+ TOKEN_STRIDE=TOKEN_STRIDE,
485
+ SCALE_DIM=SCALE_DIM,
486
+ KV_BLOCK_STRIDE=kv_cache.stride(0),
487
+ num_warps=1,
488
+ )
489
+
490
+ k_quant, scale = _reference_kv_compress_norm_rope(
491
+ state_cache,
492
+ block_table,
493
+ positions,
494
+ rms_weight,
495
+ cos_sin_cache,
496
+ compress_ratio,
497
+ overlap,
498
+ use_fp4,
499
+ rms_eps=RMS_EPS,
500
+ fp8_max=FP8_MAX,
501
+ )
502
+
503
+ if use_fp4:
504
+ for i in range(num_tokens):
505
+ blk, pos = i // kv_block_size, i % kv_block_size
506
+ val_off = pos * TOKEN_STRIDE
507
+ fp4_actual = kv_cache[blk, val_off : val_off + TOKEN_STRIDE]
508
+ assert torch.equal(k_quant[i], fp4_actual), (
509
+ f"token {i}: packed nibbles differ, "
510
+ f"{(k_quant[i] != fp4_actual).sum()} "
511
+ f"/ {TOKEN_STRIDE}"
512
+ )
513
+
514
+ scale_off = kv_block_size * TOKEN_STRIDE + pos * SCALE_DIM
515
+ scale_actual = kv_cache[blk, scale_off : scale_off + SCALE_DIM]
516
+ assert torch.equal(scale_actual, scale[i]), (
517
+ f"token {i}: ue8m0 {scale_actual.tolist()} != {scale[i].tolist()}"
518
+ )
519
+
520
+ else:
521
+ k_quant = k_quant.view(torch.uint8)
522
+ for i in range(num_tokens):
523
+ blk, pos = i // kv_block_size, i % kv_block_size
524
+ val_off = pos * TOKEN_STRIDE
525
+ assert torch.equal(
526
+ k_quant[i], kv_cache[blk, val_off : val_off + TOKEN_STRIDE]
527
+ ), f"token {i}: FP8 bytes differ"
528
+
529
+ scale_off = kv_block_size * TOKEN_STRIDE + pos * SCALE_DIM
530
+ actual_scale = kv_cache[blk, scale_off : scale_off + SCALE_DIM].view(
531
+ torch.float32
532
+ )
533
+ assert torch.equal(actual_scale, scale[i : i + 1]), (
534
+ f"token {i}: scale {actual_scale.item()} != {scale[i].item()}"
535
+ )
tests/kernels/test_fused_indexer_q_rope_quant.py
+90
-17
@@ -30,6 +30,56 @@ N_HEAD = 64
30
MAX_POS = 4096
31
32
33
+def quantize_to_mxfp4(
34
+ x: torch.Tensor,
35
+) -> tuple[torch.Tensor, torch.Tensor]:
36
+ """Reference MXFP4 quantization.
37
+
38
+ Args:
39
+ x: [..., head_dim] where head_dim is divisible by 32
40
+ Returns:
41
+ packed: [..., head_dim//2] uint8 2 E2M1 nibbles/byte, low nibble = even index
42
+ scales: [..., head_dim//32] uint8 1 ue8m0 byte
43
+ """
44
+ MXFP4_BLOCK_SIZE = 32
45
+ orig_shape = x.shape
46
+ head_dim = orig_shape[-1]
47
+ n_blocks = head_dim // MXFP4_BLOCK_SIZE
48
+
49
+ x_f32 = x.float().reshape(-1, n_blocks, MXFP4_BLOCK_SIZE)
50
+
51
+ # Per-block ue8m0 scale: 2^ceil(log2(amax / 6.0)), stored as byte = exp + 127
52
+ # 6 * 2^-126 is from https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro/blob/main/inference/kernel.py#L163
53
+ amax = x_f32.abs().amax(dim=-1, keepdim=True).clamp(min=6 * (2**-126))
54
+ log2_ratio = (amax * (1.0 / 6.0)).log2().ceil().clamp(-127.0, 127.0)
55
+ scale = log2_ratio.exp2()
56
+ ue8m0 = (log2_ratio + 127.0).to(torch.uint8) # [*, n_blocks]
57
+
58
+ # E2M1 round-to-nearest-even: midpoints round to the even code.
59
+ # E2M1 values: [0.00, 0.50, 1.00, 1.50, 2.00, 3.00, 4.00, 6.00]
60
+ # boundaries: [ 0.25, 0.75, 1.25, 1.75, 2.50, 3.50, 5.00]
61
+ x_scaled = (x_f32 / scale).clamp(-6.0, 6.0)
62
+ abs_x = x_scaled.abs()
63
+ code = torch.zeros_like(abs_x, dtype=torch.int32)
64
+ code = torch.where(abs_x > 0.25, 1, code)
65
+ code = torch.where(abs_x >= 0.75, 2, code)
66
+ code = torch.where(abs_x > 1.25, 3, code)
67
+ code = torch.where(abs_x >= 1.75, 4, code)
68
+ code = torch.where(abs_x > 2.5, 5, code)
69
+ code = torch.where(abs_x >= 3.5, 6, code)
70
+ code = torch.where(abs_x > 5.0, 7, code)
71
+ sign = ((x_scaled.view(torch.int32) >> 31) & 1).to(torch.uint8)
72
+ nibble = code.to(torch.uint8) | (sign << 3)
73
+
74
+ # Pack: even-index element → low nibble, odd-index → high nibble
75
+ nibble_flat = nibble.reshape(-1, head_dim)
76
+ packed = (nibble_flat[:, 0::2] | (nibble_flat[:, 1::2] << 4)).contiguous()
77
+ packed = packed.reshape(*orig_shape[:-1], head_dim // 2)
78
+
79
+ scales = ue8m0.view(*orig_shape[:-1], n_blocks)
80
+ return packed, scales
81
+
82
+
83
def _reference(
84
positions: torch.Tensor,
85
q: torch.Tensor,
@@ -37,6 +87,7 @@ def _reference(
87
weights: torch.Tensor,
88
softmax_scale: float,
89
head_scale: float,
90
+ use_fp4: bool = False,
91
) -> tuple[torch.Tensor, torch.Tensor]:
92
q_rot = q.clone()
93
ops.rotary_embedding(
@@ -49,22 +100,33 @@ def _reference(
100
HEAD_DIM - ROPE_DIM, # rope_dim_offset → rotate the tail
101
False,
102
)
52
- q_fp8, q_scale = per_token_group_quant_fp8(
53
- q_rot.view(-1, HEAD_DIM).contiguous(),
54
- HEAD_DIM,
55
- use_ue8m0=True,
56
- )
57
- q_fp8 = q_fp8.view(-1, N_HEAD, HEAD_DIM)
58
- q_scale = q_scale.view(-1, N_HEAD)
103
60
- weights_out = weights.to(torch.float32) * q_scale * softmax_scale * head_scale
61
- return q_fp8, weights_out
104
+ if use_fp4:
105
+ q_packed, ue8m0 = quantize_to_mxfp4(q_rot.view(-1, N_HEAD, HEAD_DIM))
106
+ # Pack 4 ue8m0 bytes into 1 int32
107
+ q_scale = ue8m0.view(torch.int32).squeeze(-1)
108
+ # FP4 path: q_scale stays separate (cannot be folded into a per-token scalar)
109
+ weights_out = weights.to(torch.float32) * softmax_scale * head_scale
110
+ return (q_packed, q_scale), weights_out
111
+
112
+ else:
113
+ q_fp8, q_scale = per_token_group_quant_fp8(
114
+ q_rot.view(-1, HEAD_DIM).contiguous(),
115
+ HEAD_DIM,
116
+ use_ue8m0=True,
117
+ )
118
+ q_fp8 = q_fp8.view(-1, N_HEAD, HEAD_DIM)
119
+ q_scale = q_scale.view(-1, N_HEAD)
120
+
121
+ weights_out = weights.to(torch.float32) * q_scale * softmax_scale * head_scale
122
+ return q_fp8, weights_out
123
124
125
@pytest.mark.parametrize("num_tokens", [1, 7, 32, 257])
126
@pytest.mark.parametrize("cache_dtype", [torch.float32, torch.bfloat16])
127
+@pytest.mark.parametrize("use_fp4", [False, True])
128
@torch.inference_mode()
67
-def test_fused_indexer_q_rope_quant_matches_unfused(num_tokens, cache_dtype):
129
+def test_fused_indexer_q_rope_quant_matches_unfused(num_tokens, cache_dtype, use_fp4):
130
device = "cuda"
131
torch.manual_seed(0)
132
@@ -77,21 +139,32 @@ def test_fused_indexer_q_rope_quant_matches_unfused(num_tokens, cache_dtype):
139
softmax_scale = HEAD_DIM**-0.5
140
head_scale = N_HEAD**-0.5
141
80
- q_fp8_ref, weights_ref = _reference(
81
- positions, q, cos_sin_cache, weights, softmax_scale, head_scale
142
+ q_quant_ref, weights_ref = _reference(
143
+ positions, q, cos_sin_cache, weights, softmax_scale, head_scale, use_fp4
144
)
83
- q_fp8_fused, weights_fused = fused_indexer_q_rope_quant(
84
- positions, q.clone(), cos_sin_cache, weights, softmax_scale, head_scale
145
+ q_quant_fused, weights_fused = fused_indexer_q_rope_quant(
146
+ positions, q.clone(), cos_sin_cache, weights, softmax_scale, head_scale, use_fp4
147
)
148
149
+ if use_fp4:
150
+ q_quant_ref, q_scale_ref = q_quant_ref
151
+ q_quant_fused, q_scale_fused = q_quant_fused
152
+
153
+ assert torch.equal(q_scale_ref, q_scale_fused), (
154
+ f"q_scale mismatch: "
155
+ f"{(q_scale_ref != q_scale_fused).sum().item()} "
156
+ f"/ {q_scale_ref.numel()} bytes differ"
157
+ )
158
+
159
# fp8 tensors aren't directly comparable via torch.equal — reinterpret as int8.
88
- ref_bits = q_fp8_ref.view(torch.int8)
89
- fused_bits = q_fp8_fused.view(torch.int8)
160
+ ref_bits = q_quant_ref.view(torch.int8)
161
+ fused_bits = q_quant_fused.view(torch.int8)
162
assert torch.equal(ref_bits, fused_bits), (
91
- f"q_fp8 mismatch: "
163
+ f"q_quant_fused mismatch: "
164
f"{(ref_bits != fused_bits).sum().item()} / {ref_bits.numel()} bytes differ"
165
)
166
167
+ assert weights_fused.dtype == torch.float32
168
assert torch.equal(weights_ref, weights_fused), (
169
f"weights mismatch: max abs diff "
170
f"{(weights_ref - weights_fused).abs().max().item()}"
vllm/v1/attention/ops/deepseek_v4_ops/fused_compress_quant_cache.py
+6
-6
@@ -21,7 +21,7 @@ and N_QUANT_BLOCKS ue8m0 bytes.
21
22
from vllm.triton_utils import tl, triton
23
24
-from .fused_indexer_q import _e2m1_nibble
24
+from .fused_indexer_q import _fp32x2_to_fp4x2
25
26
27
# =============================================================================
@@ -566,18 +566,18 @@ def _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn(
566
tl.max(tl.abs(even_2d), axis=1),
567
tl.max(tl.abs(odd_2d), axis=1),
568
)
569
- amax = tl.maximum(amax, 1e-4)
569
+ amax = tl.maximum(amax, 6.0 * (2**-126))
570
571
# ue8m0 block scale: 2^ceil(log2(amax / 6.0)), stored as (exp + 127) byte.
572
- log2_ratio = tl.ceil(tl.log2(amax / 6.0))
572
+ log2_ratio = tl.ceil(tl.log2(amax * (1.0 / 6.0)))
573
log2_ratio = tl.minimum(tl.maximum(log2_ratio, -127.0), 127.0)
574
inv_scale = tl.exp2(-log2_ratio)
575
ue8m0 = (log2_ratio + 127.0).to(tl.uint8) # [N_QUANT_BLOCKS]
576
577
inv_scale_col = tl.reshape(inv_scale, (N_QUANT_BLOCKS, 1))
578
- lo_nib = _e2m1_nibble(even_2d * inv_scale_col) # (N_BLOCKS, HALF_BLOCK) uint8
579
- hi_nib = _e2m1_nibble(odd_2d * inv_scale_col)
580
- packed = lo_nib | (hi_nib << 4)
578
+ packed = _fp32x2_to_fp4x2(
579
+ even_2d * inv_scale_col, odd_2d * inv_scale_col
580
+ ) # (N_BLOCKS, HALF_BLOCK) uint8
581
packed_flat = tl.reshape(packed, (TOKEN_STRIDE,))
582
583
tl.store(val_ptr + tl.arange(0, TOKEN_STRIDE), packed_flat)
vllm/v1/attention/ops/deepseek_v4_ops/fused_indexer_q.py
+20
-35
@@ -24,36 +24,22 @@ def _get_cos_sin(
24
25
26
@triton.jit
27
-def _e2m1_nibble(x):
28
- """Quantize fp32 x (already scale-divided) to E2M1 4-bit nibble in uint8.
29
- Matches torch.bucketize with boundaries
30
- [0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0] and right=False (each boundary
31
- belongs to the lower bucket), plus sign bit."""
32
- abs_x = tl.minimum(tl.abs(x), 6.0)
33
- code = tl.where(
34
- abs_x <= 0.25,
35
- 0.0,
36
- tl.where(
37
- abs_x <= 0.75,
38
- 1.0,
39
- tl.where(
40
- abs_x <= 1.25,
41
- 2.0,
42
- tl.where(
43
- abs_x <= 1.75,
44
- 3.0,
45
- tl.where(
46
- abs_x <= 2.5,
47
- 4.0,
48
- tl.where(abs_x <= 3.5, 5.0, tl.where(abs_x <= 5.0, 6.0, 7.0)),
49
- ),
50
- ),
51
- ),
52
- ),
53
- )
54
- code_u8 = code.to(tl.uint8)
55
- sign = ((x < 0) & (code_u8 != 0)).to(tl.uint8)
56
- return code_u8 | (sign << 3)
27
+def _fp32x2_to_fp4x2(x_lo, x_hi):
28
+ # NOTE: $1 is high nibble, $2 is low nibble
29
+ return tl.inline_asm_elementwise(
30
+ """
31
+ {
32
+ .reg .b8 tmp;
33
+ cvt.rn.satfinite.e2m1x2.f32 tmp, $1, $2;
34
+ cvt.u32.u8 $0, tmp;
35
+ }
36
+ """,
37
+ constraints="=r,f,f",
38
+ args=[x_hi, x_lo],
39
+ dtype=tl.uint32,
40
+ is_pure=True,
41
+ pack=1,
42
+ ).to(tl.uint8)
43
44
45
@triton.jit
@@ -65,17 +51,16 @@ def _quantize_mxfp4_pair(x_lo, x_hi):
51
- ue8m0 : scalar uint8 (block scale = 2^(ue8m0 - 127))
52
"""
53
amax = tl.maximum(tl.max(tl.abs(x_lo)), tl.max(tl.abs(x_hi)))
68
- amax = tl.maximum(amax, 1e-4)
54
+ # 6 * 2^-126 is from https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro/blob/main/inference/kernel.py#L163
55
+ amax = tl.maximum(amax, 6.0 * (2**-126))
56
# ue8m0 block scale: 2^ceil(log2(amax/6.0)).
70
- log2_ratio = tl.math.ceil(tl.math.log2(amax / 6.0))
57
+ log2_ratio = tl.math.ceil(tl.math.log2(amax * (1.0 / 6.0)))
58
log2_ratio = tl.minimum(tl.maximum(log2_ratio, -127.0), 127.0)
59
scale = tl.math.exp2(log2_ratio)
60
ue8m0 = (log2_ratio + 127.0).to(tl.uint8)
61
62
inv_scale = 1.0 / scale
76
- lo_nib = _e2m1_nibble(x_lo * inv_scale)
77
- hi_nib = _e2m1_nibble(x_hi * inv_scale)
78
- packed = lo_nib | (hi_nib << 4)
63
+ packed = _fp32x2_to_fp4x2(x_lo * inv_scale, x_hi * inv_scale)
64
return packed, ue8m0
65
66