V4 mxfp4

miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_mxfp8.py

@@ -94,6 +94,12 @@ def quantize_params_mxfp8(args, megatron_name, converted_named_params, quantizat
         "self_attention.linear_kv_up_proj.weight",
         "self_attention.wq_b.weight",
         "self_attention.wk.weight",
+        # DeepSeek V4 attention
+        "self_attention.wq_a.weight",
+        "self_attention.wkv.weight",
+        "self_attention.wo_b.weight",
+        "self_attention.indexer.linear_wq_b.weight",
+        "self_attention.indexer.linear_wk.weight",
     ]:
         quantize_named_params = []
         for converted_name, param in converted_named_params:

miles_plugins/models/deepseek_v4/ops/compressor.py

@@ -1,3 +1,5 @@
+import os
+
 import einops
 import torch
 import torch.nn as nn
@@ -6,7 +8,7 @@
 
 from miles_plugins.models.deepseek_v4.ops.cp_utils import all_gather_cp, get_freqs_cis_for_cp
 from miles_plugins.models.deepseek_v4.ops.kernel.precision_aligned_ops import linear_bf16_fp32
-from miles_plugins.models.deepseek_v4.ops.qat import fp8_simulate_qat
+from miles_plugins.models.deepseek_v4.ops.qat import fp8_simulate_qat, mxfp4_simulate_qat
 from miles_plugins.models.deepseek_v4.ops.rope import apply_rotary_emb, wrapped_precompute_freqs_cis
 from miles_plugins.models.deepseek_v4.ops.utils import rotate_activation
 
@@ -77,6 +79,7 @@ def __init__(
         self.rotate = rotate
         coff = 1 + self.overlap
         self.use_fp8_qat = config.fp8 is not None
+        self.use_mxfp4_qat = os.environ.get("USE_MXFP4_QAT", "0") == "1"
 
         self.cp_group = cp_group
         self.cp_size = cp_group.size() if cp_group is not None else 1
@@ -155,7 +158,9 @@ def forward_raw(self, x: torch.Tensor) -> torch.Tensor:
 
         if self.rotate:
             kv = rotate_activation(kv)
-            if self.use_fp8_qat:
+            if self.use_mxfp4_qat:
+                kv = mxfp4_simulate_qat(kv, 32)
+            elif self.use_fp8_qat:
                 kv = fp8_simulate_qat(kv, 128)
         else:
             if self.use_fp8_qat:

miles_plugins/models/deepseek_v4/ops/kernel/act_quant.py

@@ -164,3 +164,86 @@ def act_quant(
         x.copy_(y)
         return x
     return y, s
+
+
+@tilelang.jit(pass_configs=pass_configs)
+def fp4_quant_kernel(N, block_size=32, in_dtype=BF16, scale_dtype=FE8M0, inplace=False):
+    """Block-wise FP4 (E2M1) quantization with power-of-2 (UE8M0) scales.
+
+    ``inplace=True`` does a fused quant+dequant back to ``in_dtype`` (the MXFP4
+    fake-quant used for QAT). ``inplace=False`` emits packed E2M1 plus the UE8M0
+    scale. The pow2 scale is computed via the same IEEE-754 bit tricks as the FP8
+    path (``fast_round_scale``); the inner cast goes through ``FP4`` explicitly so
+    the value is rounded onto the E2M1 grid regardless of the output storage dtype.
+    """
+    M = T.symbolic("M")
+    fp4_max = 6.0
+    fp4_max_inv = 1.0 / fp4_max
+    blk_m = 32
+    group_size = block_size
+    compute_dtype = FP32
+    out_dtype = in_dtype if inplace else FP4
+
+    @T.prim_func
+    def fp4_quant_kernel_(
+        X: T.Tensor[(M, N), in_dtype],
+        Y: T.Tensor[(M, N), out_dtype],
+        S: T.Tensor[(M, T.ceildiv(N, group_size)), scale_dtype],
+    ):
+        with T.Kernel(T.ceildiv(M, blk_m), T.ceildiv(N, group_size), threads=128) as (
+            pid_m,
+            pid_n,
+        ):
+            x_shared = T.alloc_shared((blk_m, group_size), in_dtype)
+            x_local = T.alloc_fragment((blk_m, group_size), in_dtype)
+            amax_local = T.alloc_fragment((blk_m,), compute_dtype)
+            s_local = T.alloc_fragment((blk_m,), compute_dtype)
+            y_local = T.alloc_fragment((blk_m, group_size), out_dtype)
+            y_shared = T.alloc_shared((blk_m, group_size), out_dtype)
+
+            for _ in T.Pipelined(1, num_stages=2):
+                T.copy(X[pid_m * blk_m, pid_n * group_size], x_shared)
+                T.copy(x_shared, x_local)
+                T.reduce_absmax(x_local, amax_local, dim=1)
+                for i in T.Parallel(blk_m):
+                    amax_local[i] = T.max(amax_local[i], 6 * (2**-126))
+                    s_local[i] = fast_round_scale(amax_local[i], fp4_max_inv)
+                if inplace:
+                    for i, j in T.Parallel(blk_m, group_size):
+                        y_local[i, j] = T.Cast(
+                            out_dtype,
+                            T.Cast(compute_dtype, T.Cast(FP4, T.clamp(x_local[i, j] / s_local[i], -fp4_max, fp4_max)))
+                            * s_local[i],
+                        )
+                else:
+                    for i, j in T.Parallel(blk_m, group_size):
+                        y_local[i, j] = T.clamp(x_local[i, j] / s_local[i], -fp4_max, fp4_max)
+                for i in T.Parallel(blk_m):
+                    S[pid_m * blk_m + i, pid_n] = T.Cast(scale_dtype, s_local[i])
+                T.copy(y_local, y_shared)
+                T.copy(y_shared, Y[pid_m * blk_m, pid_n * group_size])
+
+    return fp4_quant_kernel_
+
+
+def fp4_act_quant(
+    x: torch.Tensor,
+    block_size: int = 32,
+    inplace: bool = False,
+) -> torch.Tensor:
+    """Block-wise FP4 (E2M1, ``1 x block_size`` UE8M0) quantization.
+
+    ``inplace=True`` does the fused quant+dequant back to BF16 (the MXFP4 fake-quant
+    used by QAT) and returns ``x``; ``inplace=False`` returns ``(packed_e2m1, scale)``.
+    """
+    N = x.size(-1)
+    assert N % block_size == 0
+    z = x.contiguous()
+    y = torch.empty_like(z) if inplace else z.new_empty(*z.shape[:-1], N // 2, dtype=torch.float4_e2m1fn_x2)
+    s = z.new_empty(*z.size()[:-1], N // block_size, dtype=torch.float8_e8m0fnu)
+    kernel = fp4_quant_kernel(N, block_size, inplace=inplace)
+    kernel(z.view(-1, N), y.view(-1, y.size(-1)), s.view(-1, N // block_size))
+    if inplace:
+        x.copy_(y)
+        return x
+    return y, s

miles_plugins/models/deepseek_v4/ops/kernel/fp4_act_quant.py

@@ -0,0 +1,9 @@
+# DEPRECATED: contents moved to act_quant.py.
+#
+# The FP4 (E2M1) quant kernel + `fp4_act_quant` wrapper now live in
+# `miles_plugins.models.deepseek_v4.ops.kernel.act_quant`. The other kernels that
+# used to live here (fp8_gemm / fp4_gemm / sparse_attn / hc_split_sinkhorn) were
+# unused dead code and have been dropped.
+#
+# This file is intentionally empty and can be deleted (`rm` was blocked in the
+# environment that performed the migration).

miles_plugins/models/deepseek_v4/ops/qat.py

@@ -2,7 +2,7 @@
 
 from tile_kernels.quant import per_token_cast_back
 
-from .kernel.act_quant import act_quant
+from .kernel.act_quant import act_quant, fp4_act_quant
 
 
 def fp8_simulate(x: torch.Tensor, block_size: int):
@@ -13,7 +13,10 @@ def fp8_simulate(x: torch.Tensor, block_size: int):
     the rest of the DeepSeek stack.
     """
     x_c = x.contiguous()
-    y, scale = act_quant(x_c, block_size, "ue8m0")
+    # Force fp32 scale storage: act_quant's auto dtype picks float8_e8m0fnu on
+    # Blackwell+DeepGEMM, but per_token_cast_back below only accepts int32/fp32
+    # scales. scale_fmt="ue8m0" still rounds the values to powers of two.
+    y, scale = act_quant(x_c, block_size, "ue8m0", scale_dtype=torch.float32)
 
     N = x_c.size(-1)
     y_flat = y.view(-1, N)
@@ -34,3 +37,27 @@ def backward(ctx, grad_kv):
 
 
 fp8_simulate_qat = DeepSeekV4LinearQATFunc.apply
+
+
+def mxfp4_simulate(x: torch.Tensor, block_size: int = 32):
+    """Simulate MXFP4 (E2M1) cast + dequant with ``1 x block_size`` UE8M0 (pow2) scaling.
+
+    Mirrors :func:`fp8_simulate` but for FP4: a single fused on-device kernel
+    (``fp4_act_quant`` with ``inplace=True``) rounds each value onto the E2M1 grid
+    and dequantizes back to ``x.dtype``. We operate on ``x.clone()`` so the op is
+    non-mutating (safe inside an autograd Function / straight-through estimator).
+    """
+    return fp4_act_quant(x.clone(), block_size, inplace=True).to(x.dtype)
+
+
+class DeepSeekV4MXFP4LinearQATFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, kv, block_size=32):
+        return mxfp4_simulate(kv, block_size)
+
+    @staticmethod
+    def backward(ctx, grad_kv):
+        return grad_kv, None
+
+
+mxfp4_simulate_qat = DeepSeekV4MXFP4LinearQATFunc.apply

miles_plugins/models/deepseek_v4/ops/v4_indexer.py

@@ -1,3 +1,5 @@
+import os
+
 import einops
 import torch
 from megatron.core import parallel_state
@@ -13,7 +15,7 @@
     _make_causal_cu_seqlens,
     batched_indexer_fwd,
 )
-from miles_plugins.models.deepseek_v4.ops.qat import fp8_simulate_qat
+from miles_plugins.models.deepseek_v4.ops.qat import fp8_simulate_qat, mxfp4_simulate_qat
 from miles_plugins.models.deepseek_v4.ops.rope import apply_rotary_emb, wrapped_precompute_freqs_cis
 from miles_plugins.models.deepseek_v4.ops.utils import rotate_activation
 
@@ -32,6 +34,7 @@ def __init__(self, config: TransformerConfig, pg_collection=None):
         self.rope_head_dim = config.qk_pos_emb_head_dim
         self.compress_ratio = 4
         self.use_fp8_qat = config.fp8 is not None
+        self.use_mxfp4_qat = os.environ.get("USE_MXFP4_QAT", "0") == "1"
 
         if pg_collection is None:
             pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=["tp", "cp"])
@@ -106,7 +109,9 @@ def forward(self, x: torch.Tensor, qr: torch.Tensor, mask=None, packed_seq_param
         q = einops.rearrange(q, "b s ... -> s b ...")
 
         q = rotate_activation(q)
-        if self.use_fp8_qat:
+        if self.use_mxfp4_qat:
+            q = mxfp4_simulate_qat(q, 32)
+        elif self.use_fp8_qat:
             q = fp8_simulate_qat(q, 128)
 
         k = self.compressor(x)
@@ -127,6 +132,10 @@ def forward(self, x: torch.Tensor, qr: torch.Tensor, mask=None, packed_seq_param
             cu_ks = cu_ks[cp_rank * seqlen : (cp_rank + 1) * seqlen]
             cu_ke = cu_ke[cp_rank * seqlen : (cp_rank + 1) * seqlen]
         index_scores = batched_indexer_fwd(q, k, weights.float(), cu_ks, cu_ke)
+        if self.use_mxfp4_qat:
+            # QAT: index scores are cast FP32 -> BF16 (and back) to match the
+            # BF16 top-k selector used at inference (~2x faster top-k).
+            index_scores = index_scores.to(torch.bfloat16).to(torch.float32)
 
         topk_count = min(self.index_topk, index_scores.size(-1))
         topk_indices = index_scores.topk(topk_count, dim=-1)[1]