feat: support GLM-5.3 Flash dynamic FP8 KV cache with regression cove…

docker/Dockerfile

@@ -5,7 +5,7 @@ ARG PYTHON_VERSION=3.10
 ARG MAMBA_VERSION=24.7.1-0
 ARG VLLM_VERSION=0.21.0
 ARG NIXL_REF=v1.2.0
-ARG FLASH_MLA_REF=47c35a7
+ARG FLASH_MLA_REF=ba89a3466e9470ad08ab39738d4e7bb66989e1e7
 ARG DEEPGEMM_REF=891d57b4db1071624b5c8fa0d1e51cb317fa709f
 ARG DEEPEP_REF=60d44037a702f651a6e18bd4aea65ed8409051c2
 ARG DEEPEP_NCCL_VERSION=2.30.4
@@ -65,7 +65,7 @@ RUN export CPATH=/usr/local/cuda/targets/x86_64-linux/include/cccl:/usr/local/cu
     git checkout ${FLASH_MLA_REF} && \
     git submodule update --init --recursive && \
     FLASH_MLA_DISABLE_SM100="$(if [ "${ENABLE_SM100}" = "1" ]; then echo 0; else echo 1; fi)" \
-    pip install --no-cache-dir .
+    pip install --no-build-isolation --no-cache-dir .
 
 RUN apt-get update && apt-get install -y libnuma-dev && rm -rf /var/lib/apt/lists/*
 

docs/CN/source/tutorial/fp8_kv_quantization.rst

@@ -12,18 +12,19 @@ FP8 KV 量化与校准指南
 功能概览
 --------
 
-LightLLM 的 FP8 KV 推理需要准备好的校准文件（``kv_cache_calib.json``），
+LightLLM 的静态 FP8 KV 推理需要准备好的校准文件（``kv_cache_calib.json``），
 并通过 ``--kv_quant_calibration_config_path`` 加载。
 你可以直接使用 ``test/advanced_config/`` 目录下已有的校准文件，
 也可以使用 `LightCompress <https://github.com/ModelTC/LightCompress>`_ 工具导出，或使用自有兼容文件。
 
 量化模式与后端对应
 ------------------
 
-LightLLM 支持两种 FP8 KV 量化模式：
+LightLLM 支持三种 FP8 KV 量化模式：
 
 - ``fp8kv_sph``: FP8 静态按 head 量化（Static Per-Head），每个 head 独立 scale，对应 ``fa3`` 后端
 - ``fp8kv_spt``: FP8 静态按 tensor 量化（Static Per-Tensor），K/V 各一个标量 scale，对应 ``flashinfer`` 后端
+- ``fp8kv_dsa``: 面向 DeepSeek-V3.2 和 GLM-5.3 Flash 稀疏 MLA，按 token 和模型对应的维度分组动态量化，无需校准文件
 
 校准文件与量化模式强相关：
 
@@ -32,6 +33,38 @@ LightLLM 支持两种 FP8 KV 量化模式：
 
 不建议混用不同模式的校准文件。
 
+GLM-5.3 Flash 动态 FP8 KV
+-------------------------
+
+``glm5_next`` 和 ``glm5_next_text`` 两种入口均支持 ``fp8kv_dsa``，计算精度使用 BF16：
+
+.. code-block:: console
+
+    $ python -m lightllm.server.api_server \
+        --model_dir /path/to/GLM-5.3-Flash --tp 4 \
+        --data_type bfloat16 --llm_kv_type fp8kv_dsa
+
+稀疏 MLA KV 使用 FlashMLA V4 格式：前 448 维按 64 维分组使用 E4M3 存储，
+最后 64 维保留真实 BF16 值，并附带七个 E8M0 scale 和一个对齐字节。
+query 保持 512 维 NoPE，不附加全零 RoPE 区域。
+包含 indexer 和 16 字节行对齐开销，每个稀疏层每 token 从 1168 字节降至 720 字节，减少约 38.4%。
+KDA 卷积、SSM 状态以及 indexer 的原始请求尾部仍使用原计算精度，因此总模型显存不会同比减少 38.4%。
+
+Prefill 直接使用本次计算的 BF16 KV。常规路径对命中的历史前缀进行分块反量化，
+BF16 gather buffer 每卡最多 2 GiB（2097152 个 token），短输入按实际长度分配。
+每块通过 FA3 的 ``only_qv`` 接口计算稀疏注意力，并使用 FP32 在线合并结果；
+gather buffer 随批内 KV 总长度增长，达到上限后保持固定。索引和注意力计算工作区仍需额外显存。
+当每卡有 8 个注意力 head、本次有 1～8 个 query token、批内 KV 总数不少于 65536 个 token，
+且布局满足优化条件时，使用直接读取 packed KV 的路径，省去 BF16 gather buffer。
+原生 MTP draft 层、CPU 前缀缓存和 P/D 传输使用同一存储布局。
+P/D 两端都需要设置 ``--llm_kv_type fp8kv_dsa``；使用
+``test/start_scripts/glm53/glm53_pd_1p1d.sh`` 时可设置环境变量 ``LLM_KV_TYPE=fp8kv_dsa``。
+传输页仍须容纳完整的 KDA/indexer 状态；量化后 token 页变小，可能需要增大 ``--pd_kv_page_size``。
+
+除原有 GLM-5.3 Flash 依赖外，此路径需要安装支持 Hopper V4 稀疏 FP8 decode 的 ``flash_mla`` 包。
+Dockerfile 固定使用最后兼容 Hopper 的提交 ``ba89a3466e9470ad08ab39738d4e7bb66989e1e7``。
+该模型不支持 ``fp8kv_sph`` 和 ``fp8kv_spt``。
+
 使用校准文件启动 FP8 推理
 -------------------------
 

docs/EN/source/tutorial/fp8_kv_quantization.rst

@@ -12,18 +12,19 @@ This chapter describes FP8 KV inference in LightLLM, including:
 Overview
 --------
 
-LightLLM FP8 KV inference requires a prepared calibration file (``kv_cache_calib.json``),
+LightLLM static FP8 KV inference requires a prepared calibration file (``kv_cache_calib.json``),
 which is loaded by ``--kv_quant_calibration_config_path``.
 You can use calibration files provided in ``test/advanced_config/``,
 export one with `LightCompress <https://github.com/ModelTC/LightCompress>`_, or use your own compatible file.
 
 Quantization Modes and Backend Mapping
 ------------------------------------------
 
-LightLLM supports two FP8 KV quantization modes:
+LightLLM supports three FP8 KV quantization modes:
 
 - ``fp8kv_sph``: FP8 Static Per-Head quantization, independent scale per head, uses ``fa3`` backend
 - ``fp8kv_spt``: FP8 Static Per-Tensor quantization, one scalar for K and one scalar for V, uses ``flashinfer`` backend
+- ``fp8kv_dsa``: dynamic per-token/group quantization for DeepSeek-V3.2 and GLM-5.3 Flash sparse MLA, with model-specific group sizes and no calibration
 
 Calibration files are mode-dependent:
 
@@ -32,6 +33,46 @@ Calibration files are mode-dependent:
 
 Avoid mixing calibration files across different modes.
 
+GLM-5.3 Flash Dynamic FP8 KV
+----------------------------
+
+Both ``glm5_next`` and ``glm5_next_text`` support ``fp8kv_dsa`` with BF16 computation:
+
+.. code-block:: console
+
+    $ python -m lightllm.server.api_server \
+        --model_dir /path/to/GLM-5.3-Flash --tp 4 \
+        --data_type bfloat16 --llm_kv_type fp8kv_dsa
+
+Sparse MLA KV uses FlashMLA's V4 format: the first 448 dimensions use E4M3 in
+groups of 64, the last 64 retain their actual BF16 values, and seven E8M0 scales
+occupy an eight-byte aligned region. Queries keep their original 512 NoPE
+dimensions without an additional zero-filled RoPE tail. Including the indexer
+and 16-byte row alignment, each sparse layer uses 720 bytes per token instead
+of 1168 bytes (38.4% less).
+KDA convolution/SSM state and the indexer's raw request tails retain their
+original computation dtypes, so this is not a 38.4% reduction in total model memory.
+
+Prefill uses newly computed BF16 KV directly. The general path dequantizes cached
+prefixes in chunks, with a BF16 gather buffer capped at 2 GiB (2097152 tokens)
+per rank; shorter inputs allocate only their actual length. Each chunk uses FA3's
+``only_qv`` sparse attention and results are merged online in FP32. The gather
+buffer grows with the batch's total KV length until it reaches the cap. Indices
+and attention workspaces still require additional memory.
+For supported layouts with 8 local attention heads, 1–8 query tokens and at least
+65536 total KV tokens in the batch, direct packed attention avoids the BF16 gather buffer.
+The same layout is used by native MTP draft layers, CPU prefix caching and P/D
+transfer. P and D must both use ``--llm_kv_type fp8kv_dsa``. For the supplied
+``test/start_scripts/glm53/glm53_pd_1p1d.sh`` launcher, set
+``LLM_KV_TYPE=fp8kv_dsa``. Transfer pages must still fit the full KDA/indexer state;
+increasing ``--pd_kv_page_size`` may be necessary because FP8 token pages are smaller.
+
+This path requires the ``flash_mla`` package with Hopper V4 sparse FP8 decode
+support, in addition to the existing GLM-5.3 Flash dependencies. The Dockerfile
+pins ``ba89a3466e9470ad08ab39738d4e7bb66989e1e7``, the last commit supporting
+Hopper. ``fp8kv_sph`` and
+``fp8kv_spt`` are not supported for this model.
+
 Start FP8 Inference with Calibration
 ------------------------------------
 

lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py

@@ -0,0 +1,228 @@
+import dataclasses
+import weakref
+
+import torch
+
+from .fp8_flashmla_sparse import NsaFlashMlaFp8SparseAttBackend, NsaFlashMlaFp8SparseDecodeAttState
+from .glm5_next import Glm5NextSparsePrefillState
+
+
+class Fp8Glm5NextSparseAttBackend(NsaFlashMlaFp8SparseAttBackend):
+    def __init__(self, model):
+        super().__init__(model=model)
+        # Match the backend's existing microbatch-separated ragged buffers.
+        self._prefill_workspaces = [{"smallq": None, "generic": None} for _ in self.ragged_mem_buffers]
+
+    def get_smallq_workspace(self, infer_state, q):
+        from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import (
+            _allocate_direct_smallq_workspace,
+            _allocate_selected_prefill_workspace,
+        )
+
+        slot = self._prefill_workspaces[infer_state.microbatch_index]
+        if slot["smallq"] is None:
+            slot["smallq"] = {
+                "direct": _allocate_direct_smallq_workspace(q[:1]),
+                "selected": _allocate_selected_prefill_workspace(q, batch_capacity=8),
+            }
+        if q.shape[0] == 1:
+            return slot["smallq"]["direct"]
+        batch = q.shape[0]
+        selected = slot["smallq"]["selected"]
+        slots = selected["compact"].shape[1]
+        return {
+            "compact": selected["compact"][:batch],
+            "lengths": selected["lengths"][:batch],
+            "page_table": selected["page_table"][:batch],
+            "cu_q": selected["cu_q"][: batch + 1],
+            "cu_k": selected["cu_k"][: batch + 1],
+            "selected_kv": selected["selected_kv"][: batch * slots],
+        }
+
+    def get_generic_prefill_workspace(self, infer_state, q, packed_kv, topk_indices):
+        from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import (
+            _allocate_generic_prefill_workspace,
+        )
+
+        slot = self._prefill_workspaces[infer_state.microbatch_index]
+        entry = slot["generic"]
+        rows, columns = topk_indices.shape
+        if (
+            entry is None
+            or entry["buffers"]["chunk_indices"].shape[0] < rows
+            or entry["buffers"]["chunk_indices"].shape[1] != columns
+            or entry["buffers"]["chunk_indices"].dtype != topk_indices.dtype
+            or entry["buffers"]["prefill_row_table"].numel() < packed_kv.shape[0]
+        ):
+            capacity = max(self.model.batch_max_tokens or 0, rows)
+            entry = {
+                "buffers": _allocate_generic_prefill_workspace(q, packed_kv, topk_indices, query_capacity=capacity),
+                "owner": None,
+            }
+            slot["generic"] = entry
+        mem_index = infer_state.mem_index
+        binding = (id(mem_index), mem_index.data_ptr(), tuple(topk_indices.shape), packed_kv.shape[0])
+        if entry["owner"] is None or entry["owner"]() is not infer_state or entry["binding"] != binding:
+            buffers = entry["buffers"]
+            # Capacity is persistent; the current batch's mapping is not.
+            buffers["prefill_row_table"].fill_(-1)
+            entry["workspace"] = {
+                "prefill_row_table": buffers["prefill_row_table"][: packed_kv.shape[0]],
+                "chunk_indices": buffers["chunk_indices"][:rows],
+                "chunk_lengths": buffers["chunk_lengths"][:rows],
+                "cu_seqlens_q": buffers["cu_seqlens_q"][: rows + 1],
+                "row_table_built": False,
+            }
+            entry["owner"] = weakref.ref(infer_state)
+            entry["binding"] = binding
+        return entry["workspace"]
+
+    def create_att_prefill_state(self, infer_state):
+        return Fp8Glm5NextSparsePrefillState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state))
+
+    def create_att_decode_state(self, infer_state):
+        config = getattr(self.model, "config", {})
+        # 2048 pooled entries plus at most three kpool remainder entries fit the 2112 rounded direct window.
+        bound = (
+            2051
+            if isinstance(config, dict) and config.get("index_topk") == 2048 and config.get("index_kpool") == 4
+            else None
+        )
+        return Fp8Glm5NextSparseDecodeState(
+            backend=self, _infer_state_weak_ref=weakref.ref(infer_state), decode_topk_bound=bound
+        )
+
+
+@dataclasses.dataclass
+class Fp8Glm5NextSparsePrefillState(Glm5NextSparsePrefillState):
+    def _smallq_workspace(self, q):
+        if self.backend is None:
+            return None
+        return self.backend.get_smallq_workspace(self.infer_state, q)
+
+    def _generic_prefill_workspace(self, q, packed_kv, topk_indices):
+        if self.backend is None:
+            return None
+        return self.backend.get_generic_prefill_workspace(self.infer_state, q, packed_kv, topk_indices)
+
+    def _nsa_prefill_att(self, q, kv, att_control):
+        from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import (
+            context_attention_fwd_fp8,
+        )
+
+        params = att_control.nsa_prefill_dict
+        # Only H64 uses the small-Q workspace paths.
+        if q.shape[1] != 64:
+            return context_attention_fwd_fp8(
+                q,
+                kv,
+                self.ragged_mem_index,
+                self.infer_state.mem_index,
+                params["prefill_cache_kv"],
+                self.ks,
+                params["topk_indices"],
+                params["softmax_scale"],
+            )
+
+        from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import (
+            _can_use_direct_smallq,
+            _can_use_selected_prefill,
+        )
+
+        direct_workspace = (
+            self._smallq_workspace(q)
+            if _can_use_direct_smallq(
+                q,
+                kv,
+                self.ragged_mem_index,
+                self.infer_state.mem_index,
+                params["prefill_cache_kv"],
+                params["topk_indices"],
+            )
+            else None
+        )
+        selected_eligible = direct_workspace is None and _can_use_selected_prefill(
+            q,
+            kv,
+            self.ragged_mem_index,
+            self.infer_state.mem_index,
+            params["prefill_cache_kv"],
+            params["topk_indices"],
+        )
+        workspace = self._smallq_workspace(q) if selected_eligible else None
+        generic_workspace = None
+        if direct_workspace is None and not selected_eligible:
+            model = getattr(self.backend, "model", None)
+            graph_guard = (
+                getattr(self.infer_state, "is_cuda_graph_capture", False)
+                or getattr(model, "prefill_graph", None) is not None
+                or hasattr(self.infer_state, "prefill_cuda_graph_exe_list")
+            )
+            if not graph_guard and not torch.cuda.is_current_stream_capturing():
+                generic_workspace = self._generic_prefill_workspace(q, kv, params["topk_indices"])
+        return context_attention_fwd_fp8(
+            q,
+            kv,
+            self.ragged_mem_index,
+            self.infer_state.mem_index,
+            params["prefill_cache_kv"],
+            self.ks,
+            params["topk_indices"],
+            params["softmax_scale"],
+            workspace,
+            direct_workspace,
+            generic_workspace,
+        )
+
+
+@dataclasses.dataclass
+class Fp8Glm5NextSparseDecodeState(NsaFlashMlaFp8SparseDecodeAttState):
+    decode_topk_bound: int = None
+
+    def _nsa_decode_att(self, q, packed_kv, att_control):
+        import flash_mla
+
+        from lightllm.models.glm5_next.triton_kernel import fp8_sparse_decode
+        from lightllm.models.glm5_next.triton_kernel.pad_query_heads import pad_query_heads
+
+        params = att_control.nsa_decode_dict
+        topk_mem_indices = params["topk_mem_indices"]
+        if topk_mem_indices.ndim == 2:
+            topk_mem_indices = topk_mem_indices.unsqueeze(1)
+        assert topk_mem_indices.shape[1] == 1, "FlashMLA sparse decode path currently expects seq_len_q == 1"
+
+        q_nope, _ = q
+        num_heads = q_nope.shape[1]
+        if fp8_sparse_decode._can_use_fp8_sparse_direct_decode(
+            q_nope,
+            packed_kv,
+            topk_mem_indices,
+            params["softmax_scale"],
+            params["kv_lora_rank"],
+            self.decode_topk_bound,
+            is_cuda_graph_capture=getattr(self.infer_state, "is_cuda_graph_capture", False),
+        ):
+            return fp8_sparse_decode._fp8_sparse_direct_decode(
+                q_nope, packed_kv, topk_mem_indices, params["softmax_scale"]
+            )
+        padded_q = pad_query_heads(
+            q_nope, is_cuda_graph_capture=getattr(self.infer_state, "is_cuda_graph_capture", False)
+        )
+        kv_cache = torch.as_strided(
+            packed_kv,
+            size=(packed_kv.shape[0], 1, 1, packed_kv.shape[-1]),
+            stride=(packed_kv.stride(0), packed_kv.shape[-1], packed_kv.shape[-1], packed_kv.stride(-1)),
+        )
+        output, _ = flash_mla.flash_mla_with_kvcache(
+            q=padded_q,
+            k_cache=kv_cache,
+            block_table=None,
+            cache_seqlens=None,
+            head_dim_v=params["kv_lora_rank"],
+            tile_scheduler_metadata=self.flashmla_sched_meta,
+            softmax_scale=params["softmax_scale"],
+            causal=False,
+            is_fp8_kvcache=True,
+            indices=topk_mem_indices,
+        )
+        return output[:, 0, :num_heads, :]

lightllm/common/kv_cache_mem_manager/__init__.py

@@ -9,6 +9,7 @@
 from .fp8_static_per_tensor_quant_mem_manager import FP8StaticPerTensorQuantMemManager
 from .qwen3next_mem_manager import Qwen3NextMemManager
 from .glm5_next_mem_manager import Glm5NextMemManager
+from .fp8_glm5_next_mem_manager import FP8Glm5NextMemManager
 
 __all__ = [
     "KvCacheAllocator",
@@ -23,4 +24,5 @@
     "FP8StaticPerTensorQuantMemManager",
     "Qwen3NextMemManager",
     "Glm5NextMemManager",
+    "FP8Glm5NextMemManager",
 ]