[Kernel] MXFP4 indexer cache for GLM-5.2 / DSA (glm_moe_dsa)

csrc/libtorch_stable/cache_kernels.cu

@@ -1,4 +1,5 @@
 #include "torch_utils.h"
+#include <cuda_bf16.h>
 #include "dispatch_utils.h"
 
 #include "../cuda_utils.h"
@@ -1458,6 +1459,63 @@ void cp_gather_and_upconvert_fp8_kv_cache(
           slot_mapping.const_data_ptr<int64_t>(), head_dim, quant_block_size, \
           cache_block_size, cache_block_stride, use_ue8m0);
 
+__device__ __forceinline__ uint8_t mxfp4_e2m1_code(float x) {
+  float ax = fabsf(x);
+  uint8_t c;
+  if (ax < 0.25f) c = 0; else if (ax < 0.75f) c = 1;
+  else if (ax < 1.25f) c = 2; else if (ax < 1.75f) c = 3;
+  else if (ax < 2.5f) c = 4; else if (ax < 3.5f) c = 5;
+  else if (ax < 5.0f) c = 6; else c = 7;
+  if (signbit(x)) c |= 0x8;
+  return c;
+}
+
+// Fused MXFP4 (E2M1 values + UE8M0 block-32 scales) indexer-K quant+insert.
+// Layout per token (fp4_bytes = head_dim/2 + head_dim/32): [E2M1 values | UE8M0 scales].
+// byte j = (e2m1(dim 2j+1)<<4) | e2m1(dim 2j) (low nibble = even index). scale byte = exp+127.
+__global__ void indexer_k_quant_and_cache_mxfp4_kernel(
+    const __nv_bfloat16* __restrict__ k,       // [num_tokens, head_dim]
+    uint8_t* __restrict__ kv_cache,            // [num_blocks, block_size, fp4_bytes]
+    const int64_t* __restrict__ slot_mapping,  // [num_tokens]
+    const int head_dim, const int block_size, const int fp4_bytes) {
+  const int t = blockIdx.x;
+  const int lane = threadIdx.x;
+  const int64_t slot = slot_mapping[t];
+  if (slot < 0) return;
+  const int d0 = lane * 4;
+  if (d0 >= head_dim) return;
+  float v[4];
+#pragma unroll
+  for (int i = 0; i < 4; i++)
+    v[i] = __bfloat162float(k[(int64_t)t * head_dim + d0 + i]);
+  float amax = 0.f;
+#pragma unroll
+  for (int i = 0; i < 4; i++) amax = fmaxf(amax, fabsf(v[i]));
+  // reduce amax within 8-lane (32-elem) block
+  for (int m = 4; m > 0; m /= 2)
+    amax = fmaxf(amax, __shfl_xor_sync(0xffffffff, amax, m));
+  float safe = fmaxf(amax, 6.0f * exp2f(-126.0f));
+  int e = (int)ceilf(log2f(safe / 6.0f));
+  e = max(-127, min(127, e));
+  float inv = 1.0f / exp2f((float)e);
+  uint16_t packed = (uint16_t)mxfp4_e2m1_code(v[0] * inv) |
+                    ((uint16_t)mxfp4_e2m1_code(v[1] * inv) << 4) |
+                    ((uint16_t)mxfp4_e2m1_code(v[2] * inv) << 8) |
+                    ((uint16_t)mxfp4_e2m1_code(v[3] * inv) << 12);
+  // BLOCK-SPLIT layout (matches fp8 indexer + cp_gather + paged reader):
+  //   [block_size * (head_dim/2) values][block_size * (head_dim/32) scales]
+  const int64_t bi = slot / block_size, bo = slot % block_size;
+  const int64_t block_base = bi * (int64_t)block_size * (int64_t)fp4_bytes;
+  uint8_t* vbase = kv_cache + block_base + bo * (int64_t)(head_dim / 2);
+  vbase[lane * 2 + 0] = (uint8_t)(packed & 0xFF);
+  vbase[lane * 2 + 1] = (uint8_t)((packed >> 8) & 0xFF);
+  if ((lane & 7) == 0) {
+    int blk = d0 / 32;
+    kv_cache[block_base + (int64_t)block_size * (head_dim / 2) +
+             bo * (int64_t)(head_dim / 32) + blk] = (uint8_t)(e + 127);
+  }
+}
+
 void indexer_k_quant_and_cache(
     torch::stable::Tensor& k,         // [num_tokens, head_dim]
     torch::stable::Tensor& kv_cache,  // [num_blocks, block_size, cache_stride]
@@ -1485,6 +1543,17 @@ void indexer_k_quant_and_cache(
       k.get_device_index());
   const cudaStream_t stream = get_current_cuda_stream();
 
+  if (scale_fmt == "mxfp4") {
+    int fp4_bytes = head_dim / 2 + head_dim / 32;
+    dim3 g(num_tokens), b(32);
+    indexer_k_quant_and_cache_mxfp4_kernel<<<g, b, 0, stream>>>(
+        reinterpret_cast<__nv_bfloat16*>(k.data_ptr()),
+        reinterpret_cast<uint8_t*>(kv_cache.data_ptr()),
+        slot_mapping.const_data_ptr<int64_t>(), head_dim, cache_block_size,
+        fp4_bytes);
+    return;
+  }
+
   static const std::string kv_cache_dtype = "fp8_e4m3";
   DISPATCH_BY_KV_CACHE_DTYPE(k.scalar_type(), kv_cache_dtype,
                              CALL_INDEXER_K_QUANT_AND_CACHE);

vllm/model_executor/layers/sparse_attn_indexer.py

@@ -380,13 +380,13 @@ def sparse_attn_indexer(
     if not skip_k_cache_insert:
         # scale_fmt can be None, but the function expects str
         assert scale_fmt is not None
-        assert not use_fp4_cache, "Unfused FP4 Insert is not supported yet"
+        # MXFP4 insert supported via the mxfp4 branch of indexer_k_quant_and_cache
         ops.indexer_k_quant_and_cache(
             k,
             kv_cache,
             slot_mapping,
             quant_block_size,
-            scale_fmt,
+            "mxfp4" if use_fp4_cache else scale_fmt,
         )
 
     # The buffer must be pre-filled with -1 (the "no token" sentinel) before the

vllm/model_executor/models/deepseek_v2.py

@@ -639,6 +639,39 @@ def get_attn_backend(self) -> AttentionBackend:
         return DeepseekV32IndexerBackend
 
 
+def _indexer_quant_q_mxfp4(
+    q: torch.Tensor, n_head: int, head_dim: int
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Quantize the DSA lightning-indexer query to MXFP4.
+
+    Produces the (packed E2M1 values, packed UE8M0 block scales) pair expected by
+    the MXFP4 indexer logits kernels, matching the FP4 indexer K-cache layout.
+    Block size is 32 (one UE8M0 scale per 32 elements). Done in PyTorch since the
+    indexer q is tiny relative to attention; the K-side quant is fused in the
+    cache-insert kernel.
+    """
+    rows = q.shape[0]
+    xb = q.float().reshape(-1, head_dim // 32, 32)
+    amax = xb.abs().amax(-1, keepdim=True).clamp(min=6 * 2**-126)
+    # UE8M0 block exponent (E8M0, bias 127)
+    l2 = (amax / 6.0).log2().ceil().clamp(-127, 127)
+    ue = (l2 + 127).to(torch.int32).squeeze(-1)
+    nblk = head_dim // 32
+    shift = torch.arange(nblk, device=q.device, dtype=torch.int32) * 8
+    q_scale = (ue << shift).sum(-1, dtype=torch.int32).view(-1, n_head)
+    # E2M1 (FP4) value encode via magnitude thresholds + sign bit
+    sc = (xb / l2.exp2()).reshape(rows, head_dim)
+    thr = torch.tensor(
+        [0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0], device=q.device
+    )
+    code = ((sc < 0).to(torch.uint8) << 3) | torch.bucketize(sc.abs(), thr).to(
+        torch.uint8
+    )
+    packed = (code[:, 0::2] | (code[:, 1::2] << 4)).to(torch.uint8)
+    q_fp4 = packed.view(-1, n_head, head_dim // 2)
+    return q_fp4, q_scale
+
+
 class Indexer(nn.Module):
     def __init__(
         self,
@@ -687,11 +720,25 @@ def __init__(
         self.quant_block_size = 128  # TODO: get from config
         self.topk_indices_buffer = topk_indices_buffer
 
+        # Optionally store the DSA indexer K-cache in MXFP4 (4-bit) instead of
+        # FP8, halving the indexer KV-cache footprint. Gated by the attention
+        # config flag `use_fp4_indexer_cache`. Reuses the existing MXFP4
+        # indexer insert/read kernels (added for DeepSeek-V4); this only wires
+        # them into the GLM / deepseek_v2 Indexer path.
+        self.use_fp4_cache = getattr(
+            vllm_config.attention_config, "use_fp4_indexer_cache", False
+        )
+
         # NOTE: (zyongye) we use fp8 naive cache,
         #       where we store value in fp8 and scale in fp32
         #       per self.quant_block_size element
         self.k_cache = DeepseekV32IndexerCache(
-            head_dim=self.head_dim + self.head_dim // self.quant_block_size * 4,
+            head_dim=(
+                # MXFP4: head_dim//2 packed e2m1 nibbles + head_dim//32 ue8m0 scales
+                self.head_dim // 2 + self.head_dim // 32
+                if self.use_fp4_cache
+                else self.head_dim + self.head_dim // self.quant_block_size * 4
+            ),
             dtype=torch.uint8,
             prefix=f"{prefix}.k_cache",
             cache_config=cache_config,
@@ -710,6 +757,7 @@ def __init__(
             self.max_model_len,
             self.max_total_seq_len,
             self.topk_indices_buffer,
+            use_fp4_cache=self.use_fp4_cache,
         )
 
         self.is_inplace_rope = is_inplace_rope
@@ -720,6 +768,7 @@ def __init__(
             and self.head_dim == 128
             and self.rope_dim == 64
             and self.scale_fmt is not None
+            and not self.use_fp4_cache
         )
 
     def forward(
@@ -802,6 +851,11 @@ def forward(
 
         # we only quant q here since k quant is fused with cache insertion
         q = q.view(-1, self.head_dim)
+        if self.use_fp4_cache:
+            q_fp4, q_scale = _indexer_quant_q_mxfp4(q, self.n_head, self.head_dim)
+            # MXFP4 indexer logits kernel requires fp32 weights
+            weights = weights.float() * self.softmax_scale * self.n_head_scale
+            return self.indexer_op(hidden_states, (q_fp4, q_scale), k, weights)
         q_fp8, q_scale = per_token_group_quant_fp8(
             q,
             self.quant_block_size,