diff --git a/docker/Dockerfile b/docker/Dockerfile index e597a7de1f..01076629a5 100644 --- a/docker/Dockerfile +++ b/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/* diff --git a/docs/CN/source/tutorial/fp8_kv_quantization.rst b/docs/CN/source/tutorial/fp8_kv_quantization.rst index ef91557b63..0beade7367 100644 --- a/docs/CN/source/tutorial/fp8_kv_quantization.rst +++ b/docs/CN/source/tutorial/fp8_kv_quantization.rst @@ -12,7 +12,7 @@ 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 `_ 工具导出，或使用自有兼容文件。 @@ -20,10 +20,11 @@ LightLLM 的 FP8 KV 推理需要准备好的校准文件（``kv_cache_calib.json 量化模式与后端对应 ------------------ -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 推理 ------------------------- diff --git a/docs/EN/source/tutorial/fp8_kv_quantization.rst b/docs/EN/source/tutorial/fp8_kv_quantization.rst index 13649915ec..a05710600e 100644 --- a/docs/EN/source/tutorial/fp8_kv_quantization.rst +++ b/docs/EN/source/tutorial/fp8_kv_quantization.rst @@ -12,7 +12,7 @@ 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 `_, or use your own compatible file. @@ -20,10 +20,11 @@ export one with `LightCompress `_, or 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 ------------------------------------ diff --git a/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py b/lightllm/common/basemodel/attention/nsa/fp8_glm5_next.py new file mode 100644 index 0000000000..8191578bb0 --- /dev/null +++ b/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, :] diff --git a/lightllm/common/kv_cache_mem_manager/__init__.py b/lightllm/common/kv_cache_mem_manager/__init__.py index b3597b4bc8..42fefbd4f7 100644 --- a/lightllm/common/kv_cache_mem_manager/__init__.py +++ b/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", ] diff --git a/lightllm/common/kv_cache_mem_manager/fp8_glm5_next_mem_manager.py b/lightllm/common/kv_cache_mem_manager/fp8_glm5_next_mem_manager.py new file mode 100644 index 0000000000..f28f400331 --- /dev/null +++ b/lightllm/common/kv_cache_mem_manager/fp8_glm5_next_mem_manager.py @@ -0,0 +1,19 @@ +from .glm5_next_mem_manager import Glm5NextMemManager +from .operator import LinearAttMemOperator + + +class FP8Glm5NextMemOperator(LinearAttMemOperator): + def copy_kv_to_mem_manager(self, layer_index, mem_index, kv): + from lightllm.models.glm5_next.triton_kernel.destindex_copy_kv_flashmla_fp8 import ( + destindex_copy_kv_flashmla_fp8, + ) + + packed_kv = self.mem_manager.get_att_input_params(layer_index) + destindex_copy_kv_flashmla_fp8(kv, mem_index, packed_kv) + + +class FP8Glm5NextMemManager(Glm5NextMemManager): + operator_class = FP8Glm5NextMemOperator + + def get_att_input_params(self, layer_index): + return self._layer_buffer(layer_index)[:, :, : self.linear_config.FP8_MLA_KV_BYTES] diff --git a/lightllm/common/kv_cache_mem_manager/glm5_next_mem_manager.py b/lightllm/common/kv_cache_mem_manager/glm5_next_mem_manager.py index 45a3be667c..74f0f26810 100644 --- a/lightllm/common/kv_cache_mem_manager/glm5_next_mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/glm5_next_mem_manager.py @@ -29,7 +29,7 @@ def get_cell_size(self): def _init_buffers(self, size, dtype, head_num, head_dim, layer_num): assert head_num == 1 - self.kv_buffer = torch.empty((layer_num, size + 1, 1, head_dim), dtype=dtype, device="cuda") + self.kv_buffer = torch.empty((layer_num, size + self.page_size, 1, head_dim), dtype=dtype, device="cuda") self._init_linear_att_buffers() def _layer_buffer(self, layer_index): diff --git a/lightllm/common/req_manager/glm5_next.py b/lightllm/common/req_manager/glm5_next.py index b4814afbf7..492dfeb73c 100644 --- a/lightllm/common/req_manager/glm5_next.py +++ b/lightllm/common/req_manager/glm5_next.py @@ -16,7 +16,8 @@ def __init__(self, max_request_num, max_sequence_length, mem_manager, linear_con super().__init__(max_request_num, max_sequence_length, mem_manager, linear_config) self.req_to_indexer_tail = LayerCache( size=max_request_num + 1, - dtype=linear_config.full_att_dtype, + # Raw keys/gates are computation state, even when packed KV is FP8. + dtype=linear_config.conv_state_dtype, shape=(linear_config.index_kpool + self.mtp_step, 2 * linear_config.index_head_dim), layer_num=linear_config.get_full_att_kv_layer_num_with_draft_model(), device="cuda", diff --git a/lightllm/common/state_cache_manager/glm5_next.py b/lightllm/common/state_cache_manager/glm5_next.py index 13195a74ac..cc26dcde8a 100644 --- a/lightllm/common/state_cache_manager/glm5_next.py +++ b/lightllm/common/state_cache_manager/glm5_next.py @@ -1,4 +1,5 @@ import dataclasses +import torch from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num, get_env_start_args from lightllm.utils.torch_dtype_utils import get_torch_dtype @@ -13,7 +14,9 @@ class Glm5NextCacheConfig(LinearAttCacheConfig): index_kpool: int = 4 index_head_dim: int = 128 - INDEX_PADDING_BYTES = 144 + INDEXER_K_BYTES = 132 + # FlashMLA V4 layout: 448 FP8 values, 64 BF16 values, seven E8M0 scales and one padding byte. + FP8_MLA_KV_BYTES = 584 @classmethod def from_model_config(cls, config, args): @@ -26,14 +29,18 @@ def from_model_config(cls, config, args): "deepseek_sparse_attention" if i % 4 == 3 else "linear_attention" for i in range(layers) ] dtype = get_torch_dtype(args.data_type) + use_fp8_kv = args.llm_kv_type == "fp8kv_dsa" + full_att_dtype = torch.uint8 if use_fp8_kv else dtype + mla_kv_bytes = cls.FP8_MLA_KV_BYTES if use_fp8_kv else config["kv_lora_rank"] * dtype.itemsize + # FlashMLA loads 16-byte vectors; keep every packed token row aligned. + packed_kv_bytes = ((mla_kv_bytes + cls.INDEXER_K_BYTES + 15) // 16) * 16 # Raw index keys and compression scores only live in request tails. - packed_dim = config["kv_lora_rank"] + cls.INDEX_PADDING_BYTES // dtype.itemsize return cls( tp_world_size=tp, full_att_all_num_kv_heads=1, - full_att_dtype=dtype, + full_att_dtype=full_att_dtype, full_att_num_kv_heads=1, - full_att_head_dim=packed_dim, + full_att_head_dim=packed_kv_bytes // full_att_dtype.itemsize, global_linear_k_heads=linear["num_heads"], global_linear_v_heads=linear["num_heads"], num_linear_k_heads=linear["num_heads"] // tp, diff --git a/lightllm/models/glm5_next/model.py b/lightllm/models/glm5_next/model.py index cc8e86613e..6c5ffa6280 100644 --- a/lightllm/models/glm5_next/model.py +++ b/lightllm/models/glm5_next/model.py @@ -8,7 +8,8 @@ from lightllm.common.basemodel import TpPartBaseModel from lightllm.common.basemodel.attention.linear.kda import KDALinearAttBackend from lightllm.common.basemodel.attention.nsa.glm5_next import Glm5NextSparseAttBackend -from lightllm.common.kv_cache_mem_manager import Glm5NextMemManager +from lightllm.common.basemodel.attention.nsa.fp8_glm5_next import Fp8Glm5NextSparseAttBackend +from lightllm.common.kv_cache_mem_manager import Glm5NextMemManager, FP8Glm5NextMemManager from lightllm.common.req_manager import Glm5NextReqManager from lightllm.common.state_cache_manager import Glm5NextCacheConfig from lightllm.distributed.communication_op import dist_group_manager @@ -41,6 +42,12 @@ def _verify_params(self): super()._verify_params() assert self.config["qk_rope_head_dim"] == 0, "GLM-5.3 Flash uses NoPE attention" assert not self.args.enable_tpsp_mix_mode, "GLM-5.3 Flash does not support TP/SP mixed mode" + if self.args.llm_kv_type not in ("None", "fp8kv_dsa"): + raise ValueError("GLM-5.3 Flash supports only --llm_kv_type None or fp8kv_dsa") + if self.args.llm_kv_type == "fp8kv_dsa" and ( + self.config["kv_lora_rank"] != 512 or self.data_type != torch.bfloat16 + ): + raise ValueError("GLM-5.3 Flash fp8kv_dsa requires kv_lora_rank=512 and --data_type bfloat16") def autotune_layers(self): return 4 @@ -68,9 +75,10 @@ def _init_req_manager(self): ) def _init_mem_manager(self): - self.mem_manager = Glm5NextMemManager( + mem_manager_class = FP8Glm5NextMemManager if self.args.llm_kv_type == "fp8kv_dsa" else Glm5NextMemManager + self.mem_manager = mem_manager_class( size=self.max_total_token_num, - dtype=self.data_type, + dtype=self.linear_config.full_att_dtype, num_kv_heads=1, head_dim=self.linear_config.full_att_head_dim, full_att_layer_num=self.linear_config.get_full_att_kv_layer_num_with_draft_model(), @@ -79,7 +87,10 @@ def _init_mem_manager(self): ) def _init_att_backend(self): - self.prefill_att_backend = Glm5NextSparseAttBackend(model=self) + backend_class = ( + Fp8Glm5NextSparseAttBackend if self.args.llm_kv_type == "fp8kv_dsa" else Glm5NextSparseAttBackend + ) + self.prefill_att_backend = backend_class(model=self) self.decode_att_backend = self.prefill_att_backend def _init_att_backend1(self): diff --git a/lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py b/lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py new file mode 100644 index 0000000000..d96ba8114e --- /dev/null +++ b/lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py @@ -0,0 +1,748 @@ +import torch +import triton +import triton.language as tl + + +# Bound the temporary BF16 gather buffer to 2 GiB (2097152 × 1 × 512 × 2 bytes). +KV_GATHER_CHUNK_SIZE = 2097152 + + +@triton.jit +def _direct_smallq_packed_kernel( + Q, + PackedKV, + RaggedMemIndex, + PrefillMemIndex, + PrefillKV, + KVStarts, + Topk, + PartialO, + PartialLse, + ragged_n, + q_stride_s: tl.constexpr, + q_stride_h: tl.constexpr, + q_stride_d: tl.constexpr, + packed_stride_s: tl.constexpr, + prefill_stride_s: tl.constexpr, + prefill_stride_d: tl.constexpr, + topk_stride_s: tl.constexpr, + topk_stride_d: tl.constexpr, + HEADS_: tl.constexpr, + BLOCK_H_: tl.constexpr, + PREFILL_N: tl.constexpr, + SPLITS_: tl.constexpr, + BLOCKS_: tl.constexpr, + SLOTS_: tl.constexpr, + softmax_scale: tl.constexpr, +): + qid = tl.program_id(0) + split = tl.program_id(1) + if HEADS_ == 8: + # H8 fits in one fixed 16-lane tile; the head mask excludes inactive lanes. + h = tl.arange(0, 16) + else: + h = tl.program_id(2) * BLOCK_H_ + tl.arange(0, BLOCK_H_) + d = tl.arange(0, 512) + hm = h < HEADS_ + q = tl.load( + Q + qid * q_stride_s + h[:, None] * q_stride_h + d[None, :] * q_stride_d, + hm[:, None], + 0.0, + ).to(tl.bfloat16) + if HEADS_ == 8: + acc = tl.zeros((16, 512), tl.float32) + m = tl.full((16,), -float("inf"), tl.float32) + denominator = tl.zeros((16,), tl.float32) + else: + acc = tl.zeros((BLOCK_H_, 512), tl.float32) + m = tl.full((BLOCK_H_,), -float("inf"), tl.float32) + denominator = tl.zeros((BLOCK_H_,), tl.float32) + start = tl.load(KVStarts + qid).to(tl.int64) + for b in tl.range(0, BLOCKS_, num_stages=1, loop_unroll_factor=1, disallow_acc_multi_buffer=True): + n = tl.arange(0, 64) + slot = split * SLOTS_ + b * 64 + n + top = tl.load( + Topk + qid * topk_stride_s + slot * topk_stride_d, + slot < 2176, + -1, + ).to(tl.int64) + ragged_pos = start + top + valid = (top >= 0) & (ragged_pos >= 0) & (ragged_pos < ragged_n) + mem = tl.load(RaggedMemIndex + ragged_pos, valid, 0).to(tl.int64) + current_row = tl.full((64,), -1, tl.int32) + for j in tl.static_range(PREFILL_N): + current_row = tl.where(valid & (mem == tl.load(PrefillMemIndex + j).to(tl.int64)), j, current_row) + cached = valid & (current_row < 0) + current = valid & (current_row >= 0) + packed_row = PackedKV + mem[:, None] * packed_stride_s + fp8 = tl.load( + packed_row.to(tl.pointer_type(tl.float8e4nv)) + d[None, :], + cached[:, None] & (d[None, :] < 448), + 0.0, + ).to(tl.float32) + scales = tl.load( + packed_row + 576 + d[None, :] // 64, + cached[:, None] & (d[None, :] < 448), + 127, + ) + fp8 = fp8 * (scales.to(tl.uint32) << 23).to(tl.float32, bitcast=True) + tail = tl.load( + (packed_row + 448).to(tl.pointer_type(tl.bfloat16)) + d[None, :] - 448, + cached[:, None] & (d[None, :] >= 448), + 0.0, + ).to(tl.float32) + cached_v = tl.where(d[None, :] < 448, fp8, tail) + current_v = tl.load( + PrefillKV + current_row[:, None] * prefill_stride_s + d[None, :] * prefill_stride_d, + current[:, None], + 0.0, + ).to(tl.float32) + kv = tl.where(current[:, None], current_v, cached_v).to(tl.bfloat16) + score = tl.dot(q, tl.trans(kv), out_dtype=tl.float32) * softmax_scale + score = tl.where(valid[None, :], score, -float("inf")) + block_m = tl.max(score, axis=1) + new_m = tl.maximum(m, block_m) + alpha = tl.where(m == -float("inf"), 0.0, tl.exp(m - new_m)) + p = tl.where(valid[None, :], tl.exp(score - new_m[:, None]), 0.0) + denominator = denominator * alpha + tl.sum(p, axis=1) + acc = acc * alpha[:, None] + tl.dot(p.to(tl.bfloat16), kv, out_dtype=tl.float32) + m = new_m + has = denominator > 0.0 + out = tl.where(has[:, None], acc / denominator[:, None], 0.0) + tl.store( + PartialO + ((qid * SPLITS_ + split) * HEADS_ + h[:, None]) * 512 + d[None, :], + out, + hm[:, None], + ) + tl.store( + PartialLse + (qid * SPLITS_ + split) * HEADS_ + h, + tl.where(has, m + tl.log(denominator), -float("inf")), + hm, + ) + + +@triton.jit +def _reduce_direct_smallq_kernel( + PartialO, + PartialLse, + Out, + HEADS_: tl.constexpr, + SPLITS_: tl.constexpr, + SPAD_: tl.constexpr, +): + qid = tl.program_id(0) + h = tl.program_id(1) + d = tl.program_id(2) * 128 + tl.arange(0, 128) + s = tl.arange(0, SPAD_) + sm = s < SPLITS_ + lse = tl.load(PartialLse + (qid * SPLITS_ + s) * HEADS_ + h, sm, -float("inf")) + maximum = tl.max(lse, axis=0) + safe_maximum = tl.where(maximum == -float("inf"), 0.0, maximum) + weight = tl.where(sm & (lse > -float("inf")), tl.exp(lse - safe_maximum), 0.0) + denominator = tl.sum(weight, axis=0) + partial = tl.load( + PartialO + ((qid * SPLITS_ + s[:, None]) * HEADS_ + h) * 512 + d[None, :], + sm[:, None], + 0.0, + ).to(tl.float32) + value = tl.sum(partial * weight[:, None], axis=0) / tl.where(denominator > 0.0, denominator, 1.0) + tl.store(Out + (qid * HEADS_ + h) * 512 + d, value) + + +def _can_use_direct_smallq(q, packed_kv, ragged_mem_index, prefill_mem_index, prefill_kv, topk_indices): + return ( + q.ndim == 3 + and 1 <= q.shape[0] <= 8 + and q.dtype == torch.bfloat16 + and (q.shape[1:] == (8, 512) or (q.shape[0] == 1 and q.shape[1:] == (64, 512))) + and packed_kv.ndim == 3 + and packed_kv.shape[1] == 1 + and packed_kv.shape[2] >= 584 + and ragged_mem_index.dtype == torch.int32 + and ragged_mem_index.numel() >= 65536 + and prefill_mem_index.dtype == torch.int32 + and prefill_mem_index.numel() == q.shape[0] + and prefill_kv.dtype == torch.bfloat16 + and prefill_kv.ndim == 3 + and prefill_kv.shape[0] == q.shape[0] + and prefill_kv.shape[1:] == (1, 512) + and topk_indices.dtype == torch.int32 + and topk_indices.ndim == 2 + and topk_indices.shape == (q.shape[0], 2176) + ) + + +def _allocate_direct_smallq_workspace(q): + qn, heads, _ = q.shape + splits = 34 if qn == 1 else 17 + return { + "partial_o": torch.empty((qn, splits, heads, q.shape[-1]), dtype=torch.float32, device=q.device), + "partial_lse": torch.empty((qn, splits, heads), dtype=torch.float32, device=q.device), + # The reducer requires a contiguous [Q, H, 512] destination. + "output": torch.empty(q.shape, dtype=q.dtype, device=q.device), + } + + +@torch.no_grad() +def _context_attention_fwd_fp8_direct_smallq( + q, + packed_kv, + ragged_mem_index, + prefill_mem_index, + prefill_kv, + kv_starts, + topk_indices, + softmax_scale, + direct_workspace=None, +): + qn = q.shape[0] + heads = q.shape[1] + # Q1 uses 34 64-key splits; H64 uses two H32 head tiles per split. + splits, blocks, slots = (34, 1, 64) if qn == 1 else (17, 2, 128) + block_h = 32 if heads == 64 else 16 + if heads == 8: + # H8 uses local buffers because its direct path has no reusable workspace. + partial_o = torch.empty((qn, splits, 8, 512), dtype=torch.float32, device=q.device) + partial_lse = torch.empty((qn, splits, 8), dtype=torch.float32, device=q.device) + output = torch.empty(q.shape, dtype=q.dtype, device=q.device) + launch_grid = (qn, splits) + else: + workspace = direct_workspace or _allocate_direct_smallq_workspace(q) + partial_o = workspace["partial_o"] + partial_lse = workspace["partial_lse"] + output = workspace["output"] + launch_grid = (qn, splits, triton.cdiv(heads, block_h)) + _direct_smallq_packed_kernel[launch_grid]( + q, + packed_kv, + ragged_mem_index, + prefill_mem_index, + prefill_kv, + kv_starts, + topk_indices, + partial_o, + partial_lse, + ragged_mem_index.numel(), + q.stride(0), + q.stride(1), + q.stride(2), + packed_kv.stride(0), + prefill_kv.stride(0), + prefill_kv.stride(2), + topk_indices.stride(0), + topk_indices.stride(1), + HEADS_=heads, + BLOCK_H_=block_h, + PREFILL_N=qn, + SPLITS_=splits, + BLOCKS_=blocks, + SLOTS_=slots, + softmax_scale=softmax_scale, + num_warps=8, + ) + _reduce_direct_smallq_kernel[(qn, heads, 4)]( + partial_o, + partial_lse, + output, + HEADS_=heads, + SPLITS_=splits, + SPAD_=triton.next_power_of_2(splits), + num_warps=4, + ) + return output + + +@triton.jit +def _build_prefill_row_table_kernel(PrefillMemIndex, RowTable, num_tokens, BLOCK: tl.constexpr): + row = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + mem_index = tl.load(PrefillMemIndex + row, row < num_tokens, 0) + tl.store(RowTable + mem_index, row, row < num_tokens) + + +@triton.jit +def _gather_prefill_kv_kernel( + PackedKV, + MemIndex, + PrefillRowTable, + PrefillKV, + Output, + num_tokens, + packed_stride_s: tl.constexpr, + prefill_stride_s: tl.constexpr, + prefill_stride_d: tl.constexpr, + BLOCK_ROWS: tl.constexpr, +): + row = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) + dim = tl.arange(0, 512) + mem_index = tl.load(MemIndex + row, row < num_tokens, 0).to(tl.int64) + prefill_row = tl.load(PrefillRowTable + mem_index, row < num_tokens, -1) + cached = (row < num_tokens) & (prefill_row < 0) + current = (row < num_tokens) & (prefill_row >= 0) + packed_row = PackedKV + mem_index[:, None] * packed_stride_s + + quantized = tl.load( + packed_row.to(tl.pointer_type(tl.float8e4nv)) + dim[None, :], + cached[:, None] & (dim[None, :] < 448), + 0.0, + ).to(tl.float32) + scales = tl.load(packed_row + 576 + dim[None, :] // 64, cached[:, None] & (dim[None, :] < 448), 127) + values = quantized * (scales.to(tl.uint32) << 23).to(tl.float32, bitcast=True) + tail = tl.load( + (packed_row + 448).to(tl.pointer_type(tl.bfloat16)) + dim[None, :] - 448, + cached[:, None] & (dim[None, :] >= 448), + 0, + ).to(tl.float32) + values = tl.where(dim[None, :] < 448, values, tail) + current_values = tl.load( + PrefillKV + prefill_row[:, None] * prefill_stride_s + dim[None, :] * prefill_stride_d, + current[:, None], + 0, + ).to(tl.float32) + tl.store( + Output + row[:, None] * 512 + dim[None, :], + tl.where(current[:, None], current_values, values), + row[:, None] < num_tokens, + ) + + +@triton.jit +def _compact_topk_indices_kernel( + TopkIndices, + KVStarts, + ChunkIndices, + ChunkLengths, + chunk_start, + chunk_end, + TOPK: tl.constexpr, + BLOCK: tl.constexpr, + COMPACT: tl.constexpr = True, +): + row = tl.program_id(0) + lane = tl.arange(0, BLOCK) + index = tl.load(TopkIndices + row * TOPK + lane, lane < TOPK, -1) + ragged_index = index + tl.load(KVStarts + row) + valid = (lane < TOPK) & (index >= 0) & (ragged_index >= chunk_start) & (ragged_index < chunk_end) + if COMPACT: + compact_index = tl.cumsum(valid.to(tl.int32)) - 1 + tl.store(ChunkIndices + row * TOPK + compact_index, ragged_index - chunk_start, valid) + tl.store(ChunkLengths + row, tl.sum(valid.to(tl.int32))) + else: + tl.store( + ChunkIndices + row * TOPK + lane, + tl.where(valid, ragged_index - chunk_start, -1), + lane < TOPK, + ) + + +@triton.jit +def _merge_att_output_kernel( + Output, + Lse, + ChunkOutput, + ChunkLse, + chunk_lse_stride_s: tl.constexpr, + chunk_lse_stride_h: tl.constexpr, + NUM_HEADS: tl.constexpr, + NUM_Q_TOKENS: tl.constexpr, + BLOCK_HEADS: tl.constexpr, +): + row = tl.program_id(0) + head = tl.program_id(1) * BLOCK_HEADS + tl.arange(0, BLOCK_HEADS) + dim = tl.arange(0, 512) + previous_lse = tl.load(Lse + head * NUM_Q_TOKENS + row, head < NUM_HEADS, -float("inf")) + chunk_lse = tl.load( + ChunkLse + row * chunk_lse_stride_s + head * chunk_lse_stride_h, head < NUM_HEADS, -float("inf") + ) + # FA3 uses natural-log LSE and returns +inf for an empty row. + previous_lse = tl.where(previous_lse == float("inf"), -float("inf"), previous_lse) + chunk_lse = tl.where(chunk_lse == float("inf"), -float("inf"), chunk_lse) + maximum = tl.maximum(previous_lse, chunk_lse) + maximum = tl.where(maximum == -float("inf"), 0, maximum) + previous_weight = tl.exp(previous_lse - maximum) + chunk_weight = tl.exp(chunk_lse - maximum) + denominator = previous_weight + chunk_weight + offset = (row * NUM_HEADS + head[:, None]) * 512 + dim[None, :] + previous = tl.load(Output + offset, (head[:, None] < NUM_HEADS) & (previous_weight[:, None] > 0), 0) + current = tl.load(ChunkOutput + offset, (head[:, None] < NUM_HEADS) & (chunk_weight[:, None] > 0), 0).to(tl.float32) + result = (previous * previous_weight[:, None] + current * chunk_weight[:, None]) / tl.where( + denominator > 0, denominator, 1 + )[:, None] + tl.store(Output + offset, result, head[:, None] < NUM_HEADS) + tl.store( + Lse + head * NUM_Q_TOKENS + row, + tl.where(denominator > 0, maximum + tl.log(denominator), -float("inf")), + head < NUM_HEADS, + ) + + +@triton.jit +def _selected_prefill_compact_kernel( + Topk, + KVStarts, + RaggedMemIndex, + CompactPhysical, + Lengths, + PageTable, + CuQ, + ragged_n: tl.constexpr, + packed_rows: tl.constexpr, + topk_stride_s: tl.constexpr, + topk_stride_d: tl.constexpr, + BATCH: tl.constexpr, + BLOCK: tl.constexpr, +): + qid = tl.program_id(0) + lane = tl.arange(0, BLOCK) + local = tl.load(Topk + qid * topk_stride_s + lane * topk_stride_d, lane < 2176, -1).to(tl.int64) + ragged_pos = tl.load(KVStarts + qid).to(tl.int64) + local + valid = (lane < 2176) & (local >= 0) & (ragged_pos >= 0) & (ragged_pos < ragged_n) + physical = tl.load(RaggedMemIndex + ragged_pos, valid, 0).to(tl.int64) + valid = valid & (physical >= 0) & (physical < packed_rows) + rank = tl.cumsum(valid.to(tl.int32), axis=0) - 1 + tl.store(CompactPhysical + qid * 2176 + rank, physical, valid) + # Every selected slot has a stable contiguous page; invalid tail pages are + # harmless because cache_seqlens limits FA3 to Lengths[qid]. + tl.store(PageTable + qid * 2176 + lane, qid * 2176 + lane, lane < 2176) + if qid == 0: + tl.store(CuQ + lane, lane.to(tl.int32), lane <= BATCH) + tl.store(Lengths + qid, tl.sum(valid.to(tl.int32), axis=0)) + + +@triton.jit +def _selected_prefill_gather_kernel( + CompactPhysical, + Lengths, + PrefillMemIndex, + PrefillKV, + PackedKV, + SelectedKV, + CuK, + packed_stride_s: tl.constexpr, + prefill_stride_s: tl.constexpr, + prefill_stride_d: tl.constexpr, + BATCH: tl.constexpr, + BATCH_PAD: tl.constexpr, +): + qid = tl.program_id(0) + row = tl.program_id(1) * 4 + tl.arange(0, 4) + length = tl.load(Lengths + qid) + live = row < length + physical = tl.load(CompactPhysical + qid * 2176 + row, live, 0).to(tl.int64) + current_row = tl.full((4,), -1, tl.int32) + for j in tl.static_range(BATCH_PAD): + current = j < BATCH + current_mem = tl.load(PrefillMemIndex + j, current, 0).to(tl.int64) + current_row = tl.where(live & current & (physical == current_mem), j, current_row) + cached = live & (current_row < 0) + current = live & (current_row >= 0) + dim = tl.arange(0, 512) + packed_row = PackedKV + physical[:, None] * packed_stride_s + group = tl.arange(0, 8) + raw_scales = tl.load(packed_row + 576 + group[None, :], cached[:, None] & (group[None, :] < 7), 127) + scales = (raw_scales.to(tl.uint32) << 23).to(tl.float32, bitcast=True) + scales = tl.broadcast_to(scales[:, :, None], (4, 8, 64)).reshape(4, 512) + fp8 = ( + tl.load( + packed_row.to(tl.pointer_type(tl.float8e4nv)) + dim[None, :], + cached[:, None] & (dim[None, :] < 448), + 0.0, + ).to(tl.float32) + * scales + ) + tail = tl.load( + (packed_row + 448).to(tl.pointer_type(tl.bfloat16)) + dim[None, :] - 448, + cached[:, None] & (dim[None, :] >= 448), + 0.0, + ).to(tl.float32) + cached_value = tl.where(dim[None, :] < 448, fp8, tail) + current_value = tl.load( + PrefillKV + current_row[:, None] * prefill_stride_s + dim[None, :] * prefill_stride_d, + current[:, None], + 0.0, + ).to(tl.float32) + value = tl.where(current[:, None], current_value, cached_value).to(tl.bfloat16) + tl.store(SelectedKV + (qid * 2176 + row[:, None]) * 512 + dim[None, :], value, row[:, None] < 2176) + if (qid == 0) & (tl.program_id(1) == 0): + x = tl.arange(0, BATCH_PAD) + lengths = tl.load(Lengths + x, x < BATCH, 0) + tl.store(CuK, 0) + tl.store(CuK + 1 + x, tl.cumsum(lengths, axis=0), x < BATCH) + + +def _can_use_selected_prefill(q, packed_kv, ragged_mem_index, prefill_mem_index, prefill_kv, topk_indices): + return ( + q.ndim == 3 + and 2 <= q.shape[0] <= 8 + and q.shape[1:] == (64, 512) + and q.dtype == torch.bfloat16 + and q.is_cuda + and q.stride(-1) == 1 + and packed_kv.ndim == 3 + and packed_kv.shape[1:] == (1, 584) + and packed_kv.dtype == torch.uint8 + and packed_kv.device == q.device + and ragged_mem_index.dtype == torch.int32 + and ragged_mem_index.device == q.device + and ragged_mem_index.numel() >= 65536 + and prefill_mem_index.dtype == torch.int32 + and prefill_mem_index.device == q.device + and prefill_mem_index.shape == (q.shape[0],) + and prefill_kv.dtype == torch.bfloat16 + and prefill_kv.device == q.device + and prefill_kv.shape == (q.shape[0], 1, 512) + and topk_indices.dtype == torch.int32 + and topk_indices.device == q.device + and topk_indices.shape == (q.shape[0], 2176) + ) + + +def _allocate_selected_prefill_workspace(q, batch_capacity=None): + batch = q.shape[0] if batch_capacity is None else batch_capacity + return { + "compact": torch.empty((batch, 2176), dtype=torch.int64, device=q.device), + "lengths": torch.empty(batch, dtype=torch.int32, device=q.device), + "page_table": torch.empty((batch, 2176), dtype=torch.int32, device=q.device), + "cu_q": torch.empty(batch + 1, dtype=torch.int32, device=q.device), + "cu_k": torch.empty(batch + 1, dtype=torch.int32, device=q.device), + "selected_kv": torch.empty((batch * 2176, 1, q.shape[-1]), dtype=q.dtype, device=q.device), + } + + +def _allocate_generic_prefill_workspace(q, packed_kv, topk_indices, query_capacity=None): + """Allocate generic prefill metadata up to the backend's planned query capacity.""" + capacity = q.shape[0] if query_capacity is None else query_capacity + return { + "prefill_row_table": torch.full((packed_kv.shape[0],), -1, dtype=torch.int32, device=q.device), + "chunk_indices": torch.empty((capacity, topk_indices.shape[1]), dtype=topk_indices.dtype, device=q.device), + "chunk_lengths": torch.empty(capacity, dtype=torch.int32, device=q.device), + "cu_seqlens_q": torch.arange(capacity + 1, dtype=torch.int32, device=q.device), + "row_table_built": False, + } + + +@torch.no_grad() +def _context_attention_fwd_fp8_selected_prefill( + q, + packed_kv, + ragged_mem_index, + prefill_mem_index, + prefill_kv, + kv_starts, + topk_indices, + softmax_scale, + selected_workspace=None, +): + from sgl_kernel.flash_attn import flash_attn_with_kvcache + + batch = q.shape[0] + workspace = selected_workspace or _allocate_selected_prefill_workspace(q) + batch_pad = triton.next_power_of_2(batch) + _selected_prefill_compact_kernel[(batch,)]( + topk_indices, + kv_starts, + ragged_mem_index, + workspace["compact"], + workspace["lengths"], + workspace["page_table"], + workspace["cu_q"], + ragged_mem_index.numel(), + packed_kv.shape[0], + topk_indices.stride(0), + topk_indices.stride(1), + BATCH=batch, + BLOCK=4096, + num_warps=8, + ) + _selected_prefill_gather_kernel[(batch, triton.cdiv(2176, 4))]( + workspace["compact"], + workspace["lengths"], + prefill_mem_index, + prefill_kv, + packed_kv, + workspace["selected_kv"], + workspace["cu_k"], + packed_kv.stride(0), + prefill_kv.stride(0), + prefill_kv.stride(2), + BATCH=batch, + BATCH_PAD=batch_pad, + num_warps=4, + ) + return flash_attn_with_kvcache( + q=None, + qv=q, + k_cache=None, + v_cache=workspace["selected_kv"].view(-1, 1, 1, 512), + page_table=workspace["page_table"], + cache_seqlens=workspace["lengths"], + cu_seqlens_q=workspace["cu_q"], + cu_seqlens_k_new=workspace["cu_k"], + max_seqlen_q=1, + softmax_scale=softmax_scale, + causal=False, + only_qv=True, + ) + + +@torch.no_grad() +def context_attention_fwd_fp8( + q: torch.Tensor, + packed_kv: torch.Tensor, + ragged_mem_index: torch.Tensor, + prefill_mem_index: torch.Tensor, + prefill_kv: torch.Tensor, + kv_starts: torch.Tensor, + topk_indices: torch.Tensor, + softmax_scale: float, + selected_workspace=None, + direct_workspace=None, + generic_workspace=None, +): + """Run NoPE sparse prefill with a bounded gather buffer and BF16 current KV.""" + # Direct packed attention is intentionally restricted to the validated + # small-Q GLM-5.3 shape; all other cases retain the bounded-gather path. + if _can_use_direct_smallq(q, packed_kv, ragged_mem_index, prefill_mem_index, prefill_kv, topk_indices): + return _context_attention_fwd_fp8_direct_smallq( + q, + packed_kv, + ragged_mem_index, + prefill_mem_index, + prefill_kv, + kv_starts, + topk_indices, + softmax_scale, + direct_workspace, + ) + if _can_use_selected_prefill(q, packed_kv, ragged_mem_index, prefill_mem_index, prefill_kv, topk_indices): + return _context_attention_fwd_fp8_selected_prefill( + q, + packed_kv, + ragged_mem_index, + prefill_mem_index, + prefill_kv, + kv_starts, + topk_indices, + softmax_scale, + selected_workspace, + ) + from sgl_kernel.flash_attn import flash_attn_with_kvcache + + num_q_tokens, num_heads, _ = q.shape + gather_chunk_size = min(KV_GATHER_CHUNK_SIZE, ragged_mem_index.numel()) + if generic_workspace is None: + # Build the row table before gathering KV and allocating compaction metadata. + prefill_row_table = torch.full((packed_kv.shape[0],), -1, dtype=torch.int32, device=q.device) + _build_prefill_row_table_kernel[(triton.cdiv(prefill_mem_index.numel(), 256),)]( + prefill_mem_index, prefill_row_table, prefill_mem_index.numel(), BLOCK=256 + ) + # Keep this potentially 2 GiB buffer local to a DSA layer. + gathered_kv = q.new_empty((gather_chunk_size, 1, 512)) + chunk_indices = torch.empty_like(topk_indices) + chunk_lengths = torch.empty(num_q_tokens, dtype=torch.int32, device=q.device) + cu_seqlens_q = torch.arange(num_q_tokens + 1, dtype=torch.int32, device=q.device) + else: + prefill_row_table = generic_workspace["prefill_row_table"] + chunk_indices = generic_workspace["chunk_indices"] + chunk_lengths = generic_workspace["chunk_lengths"] + cu_seqlens_q = generic_workspace["cu_seqlens_q"] + if not generic_workspace["row_table_built"]: + _build_prefill_row_table_kernel[(triton.cdiv(prefill_mem_index.numel(), 256),)]( + prefill_mem_index, prefill_row_table, prefill_mem_index.numel(), BLOCK=256 + ) + generic_workspace["row_table_built"] = True + # Keep this potentially 2 GiB buffer local to a DSA layer. + gathered_kv = q.new_empty((gather_chunk_size, 1, 512)) + output = lse = None + for chunk_start in range(0, ragged_mem_index.numel(), gather_chunk_size): + num_chunk_tokens = min(gather_chunk_size, ragged_mem_index.numel() - chunk_start) + _gather_prefill_kv_kernel[(triton.cdiv(num_chunk_tokens, 4),)]( + packed_kv, + ragged_mem_index[chunk_start : chunk_start + num_chunk_tokens], + prefill_row_table, + prefill_kv, + gathered_kv, + num_chunk_tokens, + packed_kv.stride(0), + prefill_kv.stride(0), + prefill_kv.stride(2), + BLOCK_ROWS=4, + ) + chunk_kv = gathered_kv[:num_chunk_tokens].view(-1, 1, 1, 512) + if ( + num_q_tokens > 8 + and num_heads == 64 + and ragged_mem_index.numel() > num_q_tokens + and ragged_mem_index.numel() <= KV_GATHER_CHUNK_SIZE + and q.dtype == torch.bfloat16 + and q.is_contiguous() + and q.stride(-1) == 1 + and packed_kv.dtype == torch.uint8 + and packed_kv.shape[1:] == (1, 584) + and topk_indices.dtype == torch.int32 + and topk_indices.shape[1] == 2176 + and topk_indices.is_contiguous() + ): + from sgl_kernel.flash_mla import flash_mla_sparse_fwd + + _compact_topk_indices_kernel[(num_q_tokens,)]( + topk_indices, + kv_starts, + chunk_indices, + chunk_lengths, + chunk_start, + chunk_start + num_chunk_tokens, + TOPK=topk_indices.shape[1], + BLOCK=triton.next_power_of_2(topk_indices.shape[1]), + COMPACT=False, + ) + return flash_mla_sparse_fwd( + q=q, + kv=gathered_kv[:num_chunk_tokens], + indices=chunk_indices.unsqueeze(1), + sm_scale=softmax_scale, + d_v=512, + )[0] + _compact_topk_indices_kernel[(num_q_tokens,)]( + topk_indices, + kv_starts, + chunk_indices, + chunk_lengths, + chunk_start, + chunk_start + num_chunk_tokens, + TOPK=topk_indices.shape[1], + BLOCK=triton.next_power_of_2(topk_indices.shape[1]), + ) + chunk_output, chunk_lse, *_ = flash_attn_with_kvcache( + # only_qv ignores Q/K; these views avoid allocating placeholder tensors. + q=q[..., :64], + k_cache=chunk_kv[..., :64], + qv=q, + v_cache=chunk_kv, + page_table=chunk_indices, + cache_seqlens=chunk_lengths, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=1, + softmax_scale=softmax_scale, + causal=False, + only_qv=True, + num_splits=1, + return_softmax_lse=True, + ) + if chunk_start == 0: + output, lse = chunk_output, chunk_lse + else: + # Keep the online reduction in FP32 across all remaining chunks. + output = output.float() + _merge_att_output_kernel[(num_q_tokens, triton.cdiv(num_heads, 4))]( + output, + lse, + chunk_output, + chunk_lse, + chunk_lse.stride(1), + chunk_lse.stride(0), + NUM_HEADS=num_heads, + NUM_Q_TOKENS=num_q_tokens, + BLOCK_HEADS=4, + ) + del chunk_output, chunk_lse + return output.to(q.dtype) diff --git a/lightllm/models/glm5_next/triton_kernel/destindex_copy_kv_flashmla_fp8.py b/lightllm/models/glm5_next/triton_kernel/destindex_copy_kv_flashmla_fp8.py new file mode 100644 index 0000000000..7a28bc6614 --- /dev/null +++ b/lightllm/models/glm5_next/triton_kernel/destindex_copy_kv_flashmla_fp8.py @@ -0,0 +1,52 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def _fwd_kernel_destindex_copy_kv_flashmla_fp8( + KV, + MemIndex, + PackedKV, + kv_stride_s, + kv_stride_d, + packed_stride_s, +): + token = tl.program_id(0) + mem_index = tl.load(MemIndex + token).to(tl.int64) + packed_row = PackedKV + mem_index * packed_stride_s + + # Eight rows share one CTA; seven quantized groups are processed in parallel. + groups = tl.arange(0, 8) + offsets = groups[:, None] * 64 + tl.arange(0, 64)[None, :] + group_mask = groups < 7 + values = tl.load(KV + token * kv_stride_s + offsets * kv_stride_d, mask=group_mask[:, None], other=0.0).to( + tl.float32 + ) + absmax = tl.max(tl.abs(values), axis=1) + scale = tl.exp2(tl.ceil(tl.log2(tl.maximum(absmax / 448.0, 1e-4)))) + quantized = tl.clamp(values / scale[:, None], -448.0, 448.0).to(tl.float8e4nv) + tl.store(packed_row + offsets, quantized.to(tl.uint8, bitcast=True), mask=group_mask[:, None]) + tl.store(packed_row + 576 + groups, (scale.to(tl.uint32, bitcast=True) >> 23).to(tl.uint8), mask=group_mask) + + tail_offsets = tl.arange(0, 64) + tail_values = tl.load(KV + token * kv_stride_s + (448 + tail_offsets) * kv_stride_d) + tl.store((packed_row + 448).to(tl.pointer_type(tl.bfloat16)) + tail_offsets, tail_values) + tl.store(packed_row + 583, 0) + + +@torch.no_grad() +def destindex_copy_kv_flashmla_fp8(kv: torch.Tensor, mem_index: torch.Tensor, packed_kv: torch.Tensor): + assert kv.shape[1:] == (1, 512) + assert packed_kv.dtype == torch.uint8 and packed_kv.shape[1] == 1 + assert packed_kv.shape[2] >= 584 and packed_kv.stride(2) == 1 + _fwd_kernel_destindex_copy_kv_flashmla_fp8[(mem_index.numel(),)]( + kv, + mem_index.contiguous(), + packed_kv, + kv.stride(0), + kv.stride(2), + packed_kv.stride(0), + num_warps=4, + num_stages=1, + ) diff --git a/lightllm/models/glm5_next/triton_kernel/fp8_sparse_decode.py b/lightllm/models/glm5_next/triton_kernel/fp8_sparse_decode.py new file mode 100644 index 0000000000..0f716ac426 --- /dev/null +++ b/lightllm/models/glm5_next/triton_kernel/fp8_sparse_decode.py @@ -0,0 +1,168 @@ +import math + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _fp8_sparse_direct_decode_kernel( + q, + packed_kv, + topk, + partial_o, + partial_lse, + q_stride_b: tl.constexpr, + q_stride_h: tl.constexpr, + q_stride_d: tl.constexpr, + packed_stride: tl.constexpr, + topk_stride_b: tl.constexpr, + topk_stride_d: tl.constexpr, + packed_rows: tl.constexpr, + softmax_scale: tl.constexpr, + HEADS: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_K: tl.constexpr, + SPLITS: tl.constexpr, +): + qid = tl.program_id(0).to(tl.int64) + split = tl.program_id(1).to(tl.int64) + if HEADS <= BLOCK_H: + h = tl.arange(0, BLOCK_H) + else: + h = tl.program_id(2) * BLOCK_H + tl.arange(0, BLOCK_H) + h = h.to(tl.int64) + d = tl.arange(0, 512).to(tl.int64) + h_mask = h < HEADS + n = tl.arange(0, BLOCK_K) + slot = split * BLOCK_K + n + indices = tl.load(topk + qid * topk_stride_b + slot * topk_stride_d, slot < 2176, -1).to(tl.int64) + valid = (indices >= 0) & (indices < packed_rows) + packed = packed_kv + indices[:, None] * packed_stride + q_values = tl.load( + q + qid * q_stride_b + h[:, None] * q_stride_h + d[None, :] * q_stride_d, + h_mask[:, None], + 0.0, + ).to(tl.bfloat16) + group = tl.arange(0, 8) + raw_scales = tl.load( + packed + 576 + group[None, :], + valid[:, None] & (group[None, :] < 7), + 127, + ) + scales = (raw_scales.to(tl.uint32) << 23).to(tl.float32, bitcast=True) + scales = tl.broadcast_to(scales[:, :, None], (BLOCK_K, 8, 64)).reshape(BLOCK_K, 512) + fp8 = ( + tl.load( + packed.to(tl.pointer_type(tl.float8e4nv)) + d[None, :], + valid[:, None] & (d[None, :] < 448), + 0.0, + ).to(tl.float32) + * scales + ) + tail = tl.load( + (packed + 448).to(tl.pointer_type(tl.bfloat16)) + d[None, :] - 448, + valid[:, None] & (d[None, :] >= 448), + 0.0, + ).to(tl.float32) + kv = tl.where(d[None, :] < 448, fp8, tail).to(tl.bfloat16) + score = tl.dot(q_values, tl.trans(kv), out_dtype=tl.float32) * softmax_scale + score = tl.where(valid[None, :], score, -float("inf")) + maximum = tl.max(score, axis=1) + safe_maximum = tl.where(maximum == -float("inf"), 0.0, maximum) + probability = tl.where(valid[None, :], tl.exp(score - safe_maximum[:, None]), 0.0) + denominator = tl.sum(probability, axis=1) + has_values = denominator > 0.0 + values = tl.dot(probability.to(tl.bfloat16), kv, out_dtype=tl.float32) + base = (qid * SPLITS + split) * HEADS + h[:, None] + tl.store( + partial_o + base * 512 + d[None, :], + tl.where(has_values[:, None], values / denominator[:, None], 0.0), + h_mask[:, None], + ) + tl.store( + partial_lse + (qid * SPLITS + split) * HEADS + h, + tl.where(has_values, maximum + tl.log(denominator), -float("inf")), + h_mask, + ) + + +def _can_use_fp8_sparse_direct_decode( + q_nope, packed_kv, topk, softmax_scale, kv_lora_rank, bound, is_cuda_graph_capture=False +): + """Validate the GLM packed format and choose the direct decode implementation.""" + # The producer guarantees valid indices occupy at most the first 2051 slots. + # The direct kernel scans 2112 slots; unknown bounds must retain the full native scan. + if bound != 2051 or kv_lora_rank != 512 or q_nope.ndim != 3: + return False + batch, heads, dim = q_nope.shape + if batch <= 0 or heads <= 0 or dim != 512: + return False + # Direct decode wins for small graph workloads; retain native eager/larger workloads. + if not (is_cuda_graph_capture and batch * triton.cdiv(heads, 16) <= 8): + return False + if not isinstance(softmax_scale, (float, int)) or not math.isfinite(softmax_scale) or softmax_scale <= 0: + return False + return ( + q_nope.dtype == torch.bfloat16 + and q_nope.is_cuda + and packed_kv.ndim == 3 + and packed_kv.shape[1] == 1 + and packed_kv.shape[2] >= 584 + and packed_kv.dtype == torch.uint8 + and packed_kv.device == q_nope.device + and packed_kv.stride(-1) == 1 + # The tail is loaded through a BF16 pointer, requiring two-byte alignment. + and packed_kv.data_ptr() % 2 == 0 + and packed_kv.stride(0) % 2 == 0 + and topk.shape == (batch, 1, 2176) + and topk.dtype in (torch.int32, torch.int64) + and topk.device == q_nope.device + ) + + +def _fp8_sparse_direct_decode(q_nope, packed_kv, topk, softmax_scale): + from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import ( + _reduce_direct_smallq_kernel, + ) + + batch = q_nope.shape[0] + heads = q_nope.shape[1] + block_h = 16 if batch * heads <= 64 else min(32, max(16, triton.next_power_of_2(heads))) + block_k = 16 if batch == 1 and heads <= 16 else 64 + splits = 2112 // block_k + partial_o = torch.empty((batch, splits, heads, 512), device=q_nope.device, dtype=torch.float32) + partial_lse = torch.empty((batch, splits, heads), device=q_nope.device, dtype=torch.float32) + out = torch.empty(q_nope.shape, device=q_nope.device, dtype=q_nope.dtype) + launch_grid = (batch, splits) if heads <= block_h else (batch, splits, triton.cdiv(heads, block_h)) + _fp8_sparse_direct_decode_kernel[launch_grid]( + q_nope, + packed_kv, + topk, + partial_o, + partial_lse, + q_nope.stride(0), + q_nope.stride(1), + q_nope.stride(2), + packed_kv.stride(0), + topk.stride(0), + topk.stride(2), + packed_kv.shape[0], + float(softmax_scale), + heads, + block_h, + block_k, + splits, + num_warps=4 if block_k == 16 else 8, + num_stages=1, + ) + _reduce_direct_smallq_kernel[(batch, heads, 4)]( + partial_o, + partial_lse, + out, + heads, + splits, + triton.next_power_of_2(splits), + num_warps=4, + ) + return out diff --git a/lightllm/models/glm5_next/triton_kernel/pad_query_heads.py b/lightllm/models/glm5_next/triton_kernel/pad_query_heads.py new file mode 100644 index 0000000000..30a68db8ed --- /dev/null +++ b/lightllm/models/glm5_next/triton_kernel/pad_query_heads.py @@ -0,0 +1,54 @@ +import triton +import triton.language as tl + + +@triton.jit +def _fused_strided_qpad_kernel( + q, + padded_q, + q_stride_b: tl.constexpr, + q_stride_h: tl.constexpr, + q_stride_d: tl.constexpr, + num_heads: tl.constexpr, + padded_num_heads: tl.constexpr, + head_dim: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + batch = tl.program_id(0) + offset = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + in_range = offset < padded_num_heads * head_dim + head = offset // head_dim + dim = offset - head * head_dim + from_q = in_range & (head < num_heads) + value = tl.load(q + batch * q_stride_b + head * q_stride_h + dim * q_stride_d, mask=from_q, other=0.0) + tl.store(padded_q + batch * padded_num_heads * head_dim + offset, value, mask=in_range) + + +def pad_query_heads(q, is_cuda_graph_capture=False): + """Pad query heads for FlashMLA while preserving the query feature dimension.""" + num_tokens, num_heads, head_dim = q.shape + head_tile = 64 # FlashMLA head-tile alignment. + padded_num_heads = triton.cdiv(num_heads, head_tile) * head_tile + if padded_num_heads != num_heads: + if is_cuda_graph_capture: + padded_q = q.new_empty((num_tokens, 1, padded_num_heads, head_dim)) + block_size = triton.next_power_of_2(2 * head_dim) + _fused_strided_qpad_kernel[(num_tokens, triton.cdiv(padded_num_heads * head_dim, block_size))]( + q, + padded_q, + q.stride(0), + q.stride(1), + q.stride(2), + num_heads, + padded_num_heads, + head_dim=head_dim, + BLOCK_SIZE=block_size, + num_warps=4, + ) + else: + padded_q = q.new_zeros((num_tokens, 1, padded_num_heads, head_dim)) + padded_q[:, 0, :num_heads] = q + else: + # H64 BMM-transposed views already satisfy FlashMLA's last-dimension stride contract. + padded_q = q.unsqueeze(1) if num_heads == head_tile and q.stride(-1) == 1 else q.unsqueeze(1).contiguous() + return padded_q diff --git a/lightllm/models/glm5_next_mtp/model.py b/lightllm/models/glm5_next_mtp/model.py index 19cd9857df..c82962fda7 100644 --- a/lightllm/models/glm5_next_mtp/model.py +++ b/lightllm/models/glm5_next_mtp/model.py @@ -2,6 +2,7 @@ from lightllm.common.basemodel import TpPartBaseModel from lightllm.common.basemodel.attention.nsa.glm5_next import Glm5NextSparseAttBackend +from lightllm.common.basemodel.attention.nsa.fp8_glm5_next import Fp8Glm5NextSparseAttBackend from lightllm.models.glm5_next.layer_infer.transformer_layer_infer import Glm5NextTransformerLayerInfer from lightllm.models.glm5_next.layer_weights.transformer_layer_weight import Glm5NextTransformerLayerWeight from lightllm.models.llama.layer_infer.post_layer_infer import LlamaPostLayerInfer @@ -66,7 +67,10 @@ def _init_mem_manager(self): self.mem_manager = self.main_model.mem_manager def _init_att_backend(self): - self.prefill_att_backend = Glm5NextSparseAttBackend(model=self) + backend_class = ( + Fp8Glm5NextSparseAttBackend if self.args.llm_kv_type == "fp8kv_dsa" else Glm5NextSparseAttBackend + ) + self.prefill_att_backend = backend_class(model=self) self.decode_att_backend = self.prefill_att_backend def autotune_layers(self): diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index c45a76f9df..ce448e818c 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -540,8 +540,8 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: quant way is static per head kv quant. fp8kv_spt: use float8_e4m3fn to store kv cache for inference, quant way is static per tensor kv quant. - fp8kv_dsa: use DeepSeek-V3.2 DSA-specific FlashMLA FP8 sparse KV cache, - intended for the deepseek_v32 model path. + fp8kv_dsa: use dynamic per-token/group FlashMLA FP8 sparse KV cache + for DeepSeek-V3.2 or GLM-5.3 Flash. fp8kv_sph and fp8kv_spt requires --kv_quant_calibration_config_path to load pre-computed FP8 scales. Note: fp8kv_spt requires flashinfer-python>=0.6.5 (default is 0.6.3, diff --git a/lightllm/server/core/objs/req.py b/lightllm/server/core/objs/req.py index 71f1c2d5ae..ad12fab498 100644 --- a/lightllm/server/core/objs/req.py +++ b/lightllm/server/core/objs/req.py @@ -265,14 +265,18 @@ def _calcu_hybrid_cpu_cache_page_len_list(self): def _fill_input_token_hash(self): self.token_hash_list = TokenHashList() self.token_hash_list.clear() - hash_values = compute_token_list_hash(self.get_prompt_ids(), get_env_start_args().cpu_cache_token_page_size) + hash_values = compute_token_list_hash( + self.get_prompt_ids_numpy(), get_env_start_args().cpu_cache_token_page_size + ) self.token_hash_list.fill(hash_values) return def _fill_hybrid_token_hash(self): self.hybrid_token_hash_list = TokenHashList() self.hybrid_token_hash_list.clear() - hash_values = compute_token_list_hash(self.get_prompt_ids(), get_env_start_args().linear_att_hash_page_size) + hash_values = compute_token_list_hash( + self.get_prompt_ids_numpy(), get_env_start_args().linear_att_hash_page_size + ) self.hybrid_token_hash_list.fill(hash_values) return diff --git a/lightllm/utils/kv_cache_utils.py b/lightllm/utils/kv_cache_utils.py index b3047b488e..4f93317d61 100644 --- a/lightllm/utils/kv_cache_utils.py +++ b/lightllm/utils/kv_cache_utils.py @@ -33,7 +33,7 @@ logger = init_logger(__name__) -def compute_token_list_hash(tokens: List[int], cpu_cache_token_page_size: int) -> List[int]: +def compute_token_list_hash(tokens: List[int] | np.ndarray, cpu_cache_token_page_size: int) -> List[int]: if len(tokens) == 0: return [] diff --git a/skills/test_model/SKILL.md b/skills/test_model/SKILL.md index b77432b6ad..6e364defb6 100644 --- a/skills/test_model/SKILL.md +++ b/skills/test_model/SKILL.md @@ -57,3 +57,24 @@ lm_eval --model local-completions \ | 离线缓存模式 | 约 20s | 因此，除非有明确理由，`skills/test_model` 下的 `lm_eval` 测试都应默认启用离线缓存变量。 + +## 按修改范围选择端到端回归 + +- 先记录修改涉及的模型、注意力后端、缓存布局和调度路径，再选择相应子 skill；无需重跑未受影响的模型。 +- GLM-5.3 Flash 动态 `fp8kv_dsa` 使用 [专用回归流程](glm5.3-flash-fp8kv-scenarios/SKILL.md)，包含完整 GSM8K、MTP、GPU/CPU 缓存与 P/D 传输。 +- 精度评测须记录数据集完整样本数、空响应数和结果；冒烟请求不能替代完整数据集评测。 + +## 请求暂停与恢复 + +涉及 KV 分配、缓存布局、hybrid 状态或 MTP 的修改，还须测试实际调度暂停与恢复。 +使用 `test/acc/test_request_pause_resume.py`，连接开启 debug 日志且 KV 容量较小的 normal 服务。 +该用例先保存串行输出，再用短请求降低调度器的输出长度估计，并发提交较长的生成请求。 +`min_new_tokens=max_new_tokens` 保证工作量；保留 `ignore_eos=False`，让调度器使用正常的长度估计。 + +验收必须同时满足： + +1. 日志中同一个压力请求先出现 `infer paused req id`，再出现 `infer recover paused req id`。 +2. 通过 `X-Request-Id` 将内部请求 ID 对应到测试请求，所有请求正常完成且 token 数量完整。 +3. 暂停恢复后的输出 token IDs 与串行参考一致；差异需要保留并排查，不能直接忽略。 +4. 未触发暂停应判失败，调整容量或负载后重测；不能用排队等待、HTTP 重试或 RL 服务暂停接口代替。 +5. MTP 或 hybrid 模型须包含相应启用配置，覆盖 KV 与运行时状态恢复。日志、输出、启动参数归档到本轮目录。 diff --git a/skills/test_model/glm5.3-flash-fp8kv-scenarios/SKILL.md b/skills/test_model/glm5.3-flash-fp8kv-scenarios/SKILL.md new file mode 100644 index 0000000000..6e45c26eac --- /dev/null +++ b/skills/test_model/glm5.3-flash-fp8kv-scenarios/SKILL.md @@ -0,0 +1,134 @@ +--- +name: test-model-glm5.3-flash-fp8kv-scenarios +description: >- + GLM-5.3 Flash dynamic fp8kv_dsa end-to-end regression: full GSM8K with matched + BF16 KV baselines for both normal and MTP configurations, + native MTP, sparse prefix reuse, CPU reload, real request pause/resume, and P/D + transfer. Supports glm5_next and glm5_next_text. No static calibration. +--- + +# GLM-5.3 Flash 动态 FP8 KV 回归 + +继承 [通用规则](../SKILL.md)。使用本机模型目录；`model_type` 可以为 `glm5_next` +或 `glm5_next_text`，计算 dtype 为 BF16，`kv_lora_rank=512`。 +本流程测试动态 `fp8kv_dsa`，不需要校准文件；BF16 KV 用作逐配置精度对照。 + +## 准备 + +确认模型目录、8 张 Hopper GPU 的占用和所选端口。每个场景独立日志目录，保存完整命令、 +server.log、评测日志、原始响应和 summary.txt。只清理本轮启动的进程。 +服务就绪须同时检查端口监听、日志没有致命错误，并用一个实际生成请求 warmup。 + +```bash +export MODEL_DIR=/mnt/mtc/models/GLM-5.3-Flash +export HF_ALLOW_CODE_EVAL=1 HF_DATASETS_OFFLINE=1 TRANSFORMERS_OFFLINE=1 HF_HUB_OFFLINE=1 +export http_proxy= https_proxy= +export no_proxy=127.0.0.1,localhost,::1 +export NO_PROXY="$no_proxy" +export LIGHTLLM_LOG_LEVEL=debug +export LOADWORKER=8 +``` + +## 场景矩阵 + +| 场景 | 配置 | 验证 | +|---|---|---| +| BF16 KV baseline | TP4，`llm_kv_type=None`，page_size=16 | 完整 GSM8K 1319 题，作为普通 FP8 的对照 | +| FP8 | TP4，`llm_kv_type=fp8kv_dsa`，page_size=16 | 完整 GSM8K 1319 题，稀疏 prefix cache | +| BF16 KV MTP baseline | BF16 KV，原生 MTP step=2，Prefill CUDA Graph | 完整 GSM8K，作为 FP8 MTP 的对照 | +| FP8 MTP | FP8，原生 MTP step=2，Prefill CUDA Graph | 完整 GSM8K、MTP 验证与接受日志 | +| 暂停恢复 | FP8 MTP；小 KV 容量 | 真实 pause→recover→完成，输出与串行参考对照 | +| FP8 CPU Cache | 小 KV 容量，hybrid 大页和 CPU cache | GPU 淘汰后真实 CPU 命中、输出一致 | +| FP8 P/D | P TP4 + D TP4，两端启用 FP8 | 四种 K-pool 尾部长度，输出与 normal FP8 对照 | + +## 普通服务与完整 GSM8K + +按场景设置 `CUDA_VISIBLE_DEVICES`、`PORT`、`NCCL_PORT`、`KV_TYPE` 和 `LOG_DIR`。 +BF16 KV baseline 使用 `KV_TYPE=None`，FP8 使用 `KV_TYPE=fp8kv_dsa`： + +```bash +python -m lightllm.server.api_server \ + --model_dir "$MODEL_DIR" --model_name glm53 --disable_vision \ + --host 127.0.0.1 --port "$PORT" --nccl_port "$NCCL_PORT" --tp 4 \ + --data_type bfloat16 --llm_kv_type "$KV_TYPE" --page_size 16 \ + --max_total_token_num 65536 --max_req_total_len 8192 \ + --running_max_req_size 32 --batch_max_tokens 4096 --chunked_prefill_size 512 \ + --graph_max_batch_size 32 --graph_max_len_in_batch 8192 \ + --linear_att_cache_size 8 --linear_att_hash_page_size 512 \ + > "$LOG_DIR/server.log" 2>&1 +``` + +MTP 场景追加 `--mtp_mode eagle_with_att --mtp_step 2 --mtp_draft_model_dir "$MODEL_DIR"`。 +完整 GSM8K 的 MTP 场景同时追加 `--enable_prefill_cudagraph --prefill_cudagraph_max_handle_token 1024`。 +服务后台启动，就绪并 warmup 后执行： + +```bash +lm_eval --model local-completions \ + --model_args "{\"model\":\"glm53\",\"base_url\":\"http://127.0.0.1:${PORT}/v1/completions\",\"tokenizer_backend\":null,\"eos_string\":\"<|endoftext|>\",\"num_concurrent\":16,\"max_gen_toks\":512,\"max_retries\":3,\"timeout\":600}" \ + --tasks gsm8k --batch_size 16 --confirm_run_unsafe_code --log_samples \ + --output_path "$LOG_DIR/gsm8k" > "$LOG_DIR/eval.log" 2>&1 + +``` + +精度验收必须满足： + +1. 普通 FP8 对比普通 BF16 KV；FP8 MTP 对比相同 MTP step 和 Prefill CUDA Graph 配置的 BF16 KV。 + 已有结果仅在以下条件全部一致且原始产物齐全时可以复用，否则补跑。不能用普通 BF16 代替 MTP baseline。 +2. 每对实验使用相同模型权重、计算 dtype、TP、GPU 型号、缓存容量、page_size、分块及 graph 参数、 + 软件版本；除 KV 类型与服务端口、日志路径等运行位置外，配置一致。BF16 指 KV dtype，模型权重量化方式不变。 +3. 数据集版本、5-shot 样例、随机种子、贪心解码参数、输出上限、stop sequences 和评测并发完全一致。 + 不使用 `--limit`；双方必须各有 1319 个唯一 doc_id，无缺失、重复或空响应。 +4. 用 `(doc_id, filter)` 配对保存的 samples，逐题核对原始输入、标准答案和对应 hash。 + 从原始响应重新提取答案并计算 strict/flexible 分数，确认逐题得分及均值与评测报告一致。 + 检查缺少最终答案、输出截断和请求失败，保留具体样本及原因。 +5. 汇总每对 BF16/FP8 的正确题数、accuracy、差值（百分点），以及仅 BF16 正确、仅 FP8 正确的题数。 + 同时保留双方完整命令、日志、results 和 samples。未指定允许的精度下降阈值时报告实测差值， + 不能仅凭评测进程成功结束宣称精度无损或通过精度门槛。 + +以下功能场景单独验证缓存和运行时状态，不替代上述数据集精度对照，也不要求另跑 BF16 功能场景。 +在普通 FP8 和 FP8 MTP 服务上执行缓存用例，覆盖 2304～2307 token， +超过 index_topk=2048 并包含四种 K-pool 尾部长度： + +```bash +python test/acc/test_glm53_fp8kv.py \ + --url "http://127.0.0.1:$PORT" --model-dir "$MODEL_DIR" \ + --output "$LOG_DIR/cache.json" +``` + +## 暂停恢复与 CPU 回载 + +独立重启小容量 normal 服务：将上面的 `max_total_token_num` 改为 4096， +`max_req_total_len=3072`、`running_max_req_size=8`、`graph_max_batch_size=8`、 +`batch_max_tokens=1024`、`chunked_prefill_size=257`、`linear_att_cache_size=8`、 +`linear_att_hash_page_size=128`，增加 `--linear_att_page_block_num 2 --router_token_ratio 1.0`。 +使用 FP8 MTP 配置，执行公共规则要求的暂停恢复用例。 + +```bash +python test/acc/test_request_pause_resume.py \ + --url "http://127.0.0.1:$PORT" --model-dir "$MODEL_DIR" \ + --server-log "$LOG_DIR/server.log" --input-tokens 257 --output-tokens 512 \ + --concurrency 8 --output "$LOG_DIR/pause-resume.json" +``` + +FP8 CPU 场景再增加 `--enable_cpu_cache --cpu_cache_storage_size 16 --cache_placement_strategy legacy`。 +CPU 测试的淘汰请求总长度应大于 GPU KV 容量，且小于 CPU 容量： + +```bash +python test/acc/test_glm53_fp8kv.py \ + --url "http://127.0.0.1:$PORT" --model-dir "$MODEL_DIR" \ + --cpu-cache --server-log "$LOG_DIR/server.log" --output "$LOG_DIR/cpu-cache.json" +``` + +必须在目标回载请求的完成日志中看到 `cpu_prompt_cache_len > 0`，仅配置 CPU cache 不算覆盖。 + +## P/D + +使用 `test/start_scripts/glm53/glm53_pd_1p1d.sh` 的 P/D 配置，设置 +`LLM_KV_TYPE=fp8kv_dsa`，给 P、D、master 分别保存日志。两端参数须一致； +FP8 transfer page 使用 `pd_kv_page_size=32768`，16384-token 的 FP8 页不足以容纳完整 hybrid 状态。 +P/D 节点使用实际网卡 IP；master、P、D 的 `max_req_total_len` 必须一致。 +与 normal FP8 对照时保持 page_size、计算 dtype 和 MTP 配置一致。 + +通过 master 运行 `test_glm53_fp8kv.py`，并增加 +`--reference `。检查 P/D 日志证明真实传输，保留结果。 +每个场景完成后停止所属服务，最终确认 GPU 和端口释放。 diff --git a/test/acc/test_glm53_fp8kv.py b/test/acc/test_glm53_fp8kv.py new file mode 100644 index 0000000000..054b7d5808 --- /dev/null +++ b/test/acc/test_glm53_fp8kv.py @@ -0,0 +1,116 @@ +"""GLM-5.3 Flash HTTP regression for sparse prefix reuse and CPU/P-D cache transport.""" + +import argparse +import json +import re +import time +import uuid +from pathlib import Path + +import requests + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--url", required=True) + parser.add_argument("--model-dir", required=True) + parser.add_argument("--output", required=True, type=Path) + parser.add_argument("--server-log", type=Path) + parser.add_argument("--reference", type=Path, help="Compare cold outputs against a normal-mode run") + parser.add_argument("--cpu-cache", action="store_true", help="Require a CPU hit after GPU eviction") + parser.add_argument("--evict-prompts", type=int, default=4) + args = parser.parse_args() + if args.cpu_cache and args.server_log is None: + parser.error("--cpu-cache requires --server-log") + from lightllm.server.tokenizer import get_tokenizer + + tokenizer = get_tokenizer(args.model_dir) + content = "Question: What is 2 + 3?\nAnswer: 5.\n" + suffix = "Question: What is 2 + 3?\nAnswer:" + run_id = uuid.uuid4().hex + + def build_prompt(num_tokens, tag): + prefix = f"Document {tag}.\n" + repeat = (num_tokens - len(tokenizer.encode(prefix + suffix))) // len(tokenizer.encode(content)) - 1 + body = (prefix + content * repeat).rstrip() + " " + padding = num_tokens - len(tokenizer.encode(body + suffix)) + text = body + "a " * padding + suffix + assert len(tokenizer.encode(text)) == num_tokens + return text + + def generate(prompt_text, label, output_tokens=32): + input_ids = tokenizer.encode(prompt_text) + request_id = f"{run_id}-{label}" + response = requests.post( + args.url.rstrip("/") + "/generate", + headers={"X-Request-Id": request_id}, + json={ + "inputs": prompt_text, + "parameters": { + "do_sample": False, + "ignore_eos": True, + "max_new_tokens": output_tokens, + "return_details": True, + }, + }, + timeout=600, + ) + response.raise_for_status() + result = response.json() + assert result["count_output_tokens"] == output_tokens, result + assert result["finish_reason"] == "length", result + assert len(result["tokens"]) == output_tokens, result + assert result["prompt_tokens"] == len(input_ids), result + return { + "request_id": request_id, + "input_tokens": len(input_ids), + "input_ids": input_ids, + "token_ids": [t["id"] for t in result["tokens"]], + "text": result["generated_text"][0], + "cache_hit_tokens": result["tokens"][0].get("prompt_cache_len", 0), + } + + report = {"url": args.url, "cpu_cache": args.cpu_cache, "results": []} + reference = json.loads(args.reference.read_text()) if args.reference else None + try: + for remainder in range(4): + prompt_text = build_prompt(2304 + remainder, f"target-{remainder}") + row = {"remainder": remainder} + report["results"].append(row) + row["cold"] = generate(prompt_text, f"{remainder}-cold") + row["warm"] = generate(prompt_text, f"{remainder}-warm") + assert row["warm"]["cache_hit_tokens"] >= 2048, row + assert row["warm"]["token_ids"] == row["cold"]["token_ids"], row + if reference: + assert row["cold"]["input_ids"] == reference["results"][remainder]["cold"]["input_ids"], row + assert row["cold"]["token_ids"] == reference["results"][remainder]["cold"]["token_ids"], row + if args.cpu_cache: + for index in range(args.evict_prompts): + generate(build_prompt(2304, f"evict-{run_id}-{remainder}-{index}"), f"evict-{remainder}-{index}", 1) + row["reloaded"] = generate(prompt_text, f"{remainder}-reloaded") + deadline = time.monotonic() + 10 + while True: + lines = args.server_log.read_text().splitlines() + completed = [line for line in lines if f"X-Request-Id:{row['reloaded']['request_id']} " in line] + if completed: + break + if time.monotonic() >= deadline: + raise AssertionError("Missing CPU reload completion log") + time.sleep(0.2) + match = re.search(r"cpu_prompt_cache_len:(\d+)", completed[-1]) + assert match and int(match[1]) > 0, completed[-1] + row["cpu_hit_tokens"] = int(match[1]) + assert row["reloaded"]["token_ids"] == row["warm"]["token_ids"], row + report["status"] = "passed" + except Exception as exc: + report["status"] = "failed" + report["error"] = repr(exc) + raise + finally: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2)) + print(json.dumps(report), flush=True) + + +if __name__ == "__main__": + main() diff --git a/test/acc/test_request_pause_resume.py b/test/acc/test_request_pause_resume.py new file mode 100644 index 0000000000..09d164852c --- /dev/null +++ b/test/acc/test_request_pause_resume.py @@ -0,0 +1,159 @@ +"""Exercise real scheduler preemption through the streaming HTTP API.""" + +import argparse +import concurrent.futures +import json +import re +import threading +import time +import uuid +from pathlib import Path + +import requests + + +def generate_stream(url, prompt, output_tokens, request_id, timeout): + token_ids = [] + finish_reason = None + with requests.post( + url.rstrip("/") + "/generate_stream", + headers={"X-Request-Id": request_id}, + json={ + "inputs": prompt, + "parameters": { + "do_sample": False, + # Keep EOS enabled so admission uses its normal length estimate. + # min_new_tokens ensures the actual workload exceeds that estimate. + "ignore_eos": False, + "min_new_tokens": output_tokens, + "max_new_tokens": output_tokens, + }, + }, + stream=True, + timeout=timeout, + ) as response: + response.raise_for_status() + for line in response.iter_lines(chunk_size=1): + if not line.startswith(b"data:"): + continue + event = json.loads(line[5:]) + if "token" not in event: + raise RuntimeError(f"Unexpected stream event: {event}") + token_ids.append(event["token"]["id"]) + if event.get("finished"): + finish_reason = event["finish_reason"] + if len(token_ids) != output_tokens or finish_reason not in ("length", "stop"): + raise AssertionError(f"{request_id}: received {len(token_ids)} tokens, finish_reason={finish_reason}") + return {"request_id": request_id, "token_ids": token_ids, "finish_reason": finish_reason} + + +def check_pause_resume_events(log_text, request_ids): + completed_requests = {} + for line in log_text.splitlines(): + match = re.search(r"X-Request-Id:(\S+).*lightllm_req_id:(\d+)", line) + if match and match[1] in request_ids: + completed_requests[match[2]] = match[1] + paused, recovered, pending = set(), set(), set() + events = [] + for match in re.finditer(r"infer (paused|recover paused) req id (\d+)", log_text): + action, req_id = match.groups() + if req_id not in completed_requests: + continue + events.append({"action": action, "req_id": req_id, "request_id": completed_requests[req_id]}) + if action == "paused": + paused.add(req_id) + pending.add(req_id) + elif req_id in pending: + recovered.add(req_id) + pending.remove(req_id) + if len(completed_requests) != len(request_ids): + raise AssertionError("Missing completion log for a pressure request") + if not paused or paused != recovered or pending: + raise AssertionError(f"Actual pause/resume not proven: paused={sorted(paused)}, recovered={sorted(recovered)}") + return events + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--url", required=True) + parser.add_argument("--model-dir", required=True) + parser.add_argument("--server-log", required=True, type=Path) + parser.add_argument("--output", required=True, type=Path) + parser.add_argument("--input-tokens", type=int, default=512) + parser.add_argument("--output-tokens", type=int, default=512) + parser.add_argument("--concurrency", type=int, default=8) + parser.add_argument("--warmup-requests", type=int, default=64) + parser.add_argument("--timeout", type=float, default=600) + args = parser.parse_args() + from lightllm.server.tokenizer import get_tokenizer + + tokenizer = get_tokenizer(args.model_dir) + run_id = uuid.uuid4().hex + content_ids = tokenizer.encode("Continue the numbered explanation of language model inference and caching. ") + + def build_prompt(index): + prefix_ids = tokenizer.encode(f"{run_id}-{index}: ", add_special_tokens=False) + return (prefix_ids + content_ids * (args.input_tokens // len(content_ids) + 1))[: args.input_tokens] + + report = {"settings": {k: str(v) if isinstance(v, Path) else v for k, v in vars(args).items()}} + try: + # Obtain the reference before preemption, one request at a time. + prompts = [build_prompt(i) for i in range(args.concurrency)] + reference = [ + generate_stream(args.url, p, args.output_tokens, f"{run_id}-reference-{i}", args.timeout) + for i, p in enumerate(prompts) + ] + report["reference"] = reference + # Short completions lower the scheduler's EMA; distinct prompts evict + # reference KV from the deliberately small GPU cache. + for i in range(args.warmup_requests): + generate_stream(args.url, build_prompt(i + args.concurrency), 1, f"{run_id}-warmup-{i}", args.timeout) + offset = args.server_log.stat().st_size + barrier = threading.Barrier(args.concurrency) + + def run_pressure_request(index): + barrier.wait(timeout=args.timeout) + return generate_stream( + args.url, prompts[index], args.output_tokens, f"{run_id}-pressure-{index}", args.timeout + ) + + with concurrent.futures.ThreadPoolExecutor(max_workers=args.concurrency) as pool: + report["pressure"] = list(pool.map(run_pressure_request, range(args.concurrency))) + # HTTP completion is emitted just before its server-side summary log. + deadline = time.monotonic() + 10 + while True: + with args.server_log.open() as f: + f.seek(offset) + log_text = f.read() + try: + report["pause_events"] = check_pause_resume_events( + log_text, {r["request_id"] for r in report["pressure"]} + ) + break + except AssertionError: + if time.monotonic() >= deadline: + raise + time.sleep(0.2) + mismatches = [] + for expected, actual in zip(reference, report["pressure"]): + if actual["token_ids"] != expected["token_ids"]: + first_mismatch = next( + i for i, (a, b) in enumerate(zip(actual["token_ids"], expected["token_ids"])) if a != b + ) + mismatches.append({"request_id": actual["request_id"], "first_mismatch": first_mismatch}) + report["mismatches"] = mismatches + if mismatches: + raise AssertionError(f"Paused workload differs from serial reference: {mismatches}") + report["status"] = "passed" + except Exception as exc: + report["status"] = "failed" + report["error"] = repr(exc) + raise + finally: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2)) + print(json.dumps({k: v for k, v in report.items() if k not in ("reference", "pressure")}), flush=True) + + +if __name__ == "__main__": + main() diff --git a/test/start_scripts/glm53/glm53_pd_1p1d.sh b/test/start_scripts/glm53/glm53_pd_1p1d.sh index 4ad924dd74..ceab900d2a 100644 --- a/test/start_scripts/glm53/glm53_pd_1p1d.sh +++ b/test/start_scripts/glm53/glm53_pd_1p1d.sh @@ -9,6 +9,12 @@ fi PORT="$1" NODE_IP="$2" MODEL_DIR="${3:-/nvme/models/GLM-5.3-Flash}" +# Set LLM_KV_TYPE=fp8kv_dsa to quantize sparse MLA KV on both P and D. +LLM_KV_TYPE="${LLM_KV_TYPE:-None}" +PD_KV_PAGE_SIZE=16384 +if [[ "${LLM_KV_TYPE}" == "fp8kv_dsa" ]]; then + PD_KV_PAGE_SIZE=32768 +fi # P/D must advertise an address reachable by the master and the other node. export NO_PROXY="${NO_PROXY:+${NO_PROXY},}127.0.0.1,localhost,${NODE_IP}" @@ -18,6 +24,7 @@ export LOADWORKER="${LOADWORKER:-8}" COMMON_ARGS=( --model_dir "${MODEL_DIR}" --model_name glm53 + --llm_kv_type "${LLM_KV_TYPE}" --tp 4 --batch_max_tokens 8192 --running_max_req_size 64 @@ -28,7 +35,7 @@ COMMON_ARGS=( --linear_att_ssm_data_type float32 --pd_trans_mode nccl # One transfer page must also fit the global Conv/SSM/indexer-tail state. - --pd_kv_page_size 16384 + --pd_kv_page_size "${PD_KV_PAGE_SIZE}" --pd_kv_page_num 2 --pd_master_ip 127.0.0.1 --pd_master_port "${PORT}" diff --git a/unit_tests/models/glm5_next/test_cache.py b/unit_tests/models/glm5_next/test_cache.py index ae40702d75..d9095f74be 100644 --- a/unit_tests/models/glm5_next/test_cache.py +++ b/unit_tests/models/glm5_next/test_cache.py @@ -4,7 +4,7 @@ import pytest import torch -from lightllm.common.kv_cache_mem_manager import Glm5NextMemManager +from lightllm.common.kv_cache_mem_manager import Glm5NextMemManager, FP8Glm5NextMemManager from lightllm.common.req_manager import Glm5NextReqManager from lightllm.common.state_cache_manager import Glm5NextCacheConfig from lightllm.server.core.objs.start_args_type import StartArgs @@ -14,36 +14,56 @@ pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_nope_cache_config_uses_native_mla_width(): - config = Glm5NextCacheConfig.from_model_config( - { - "kv_lora_rank": 512, - "num_hidden_layers": 4, - "layer_types": ["linear_attention"] * 3 + ["deepseek_sparse_attention"], - "index_kpool": 4, - "index_head_dim": 128, - "linear_attn_config": { - "num_heads": 8, - "head_dim": 128, - "short_conv_kernel_size": 4, - "kda_layers": [0, 1, 2], - }, +@pytest.fixture +def model_config(): + return { + "kv_lora_rank": 512, + "num_hidden_layers": 4, + "layer_types": ["linear_attention"] * 3 + ["deepseek_sparse_attention"], + "index_kpool": 4, + "index_head_dim": 128, + "linear_attn_config": { + "num_heads": 8, + "head_dim": 128, + "short_conv_kernel_size": 4, + "kda_layers": [0, 1, 2], }, - StartArgs(tp=4, data_type="bfloat16"), - ) - assert config.full_att_head_dim == 584 - assert config.full_att_head_dim * config.full_att_dtype.itemsize == 512 * 2 + 144 + } + + +@pytest.mark.parametrize("kv_type", ["None", "fp8kv_dsa"]) +@pytest.mark.parametrize("page_size", [1, 16]) +def test_nope_cache_layout_and_hold_page(monkeypatch, model_config, kv_type, page_size): + args = StartArgs(tp=4, data_type="bfloat16", llm_kv_type=kv_type) + config = Glm5NextCacheConfig.from_model_config(model_config, args) + nested_config = Glm5NextCacheConfig.from_model_config({"text_config": model_config}, args) + assert nested_config == config + # The allocator reserves an entire HOLD page, including when KV is quantized. + monkeypatch.setattr(Glm5NextMemManager, "_init_linear_att_buffers", lambda self: None) + mem_manager_class = FP8Glm5NextMemManager if kv_type == "fp8kv_dsa" else Glm5NextMemManager + mem = object.__new__(mem_manager_class) + mem.page_size = page_size + mem._init_buffers(32, config.full_att_dtype, 1, config.full_att_head_dim, 1) + assert mem.kv_buffer.shape[1] == 32 + page_size + + assert config.full_att_head_dim == (720 if kv_type == "fp8kv_dsa" else 584) + assert config.full_att_dtype == (torch.uint8 if kv_type == "fp8kv_dsa" else torch.bfloat16) + assert config.conv_state_dtype == torch.bfloat16 @pytest.mark.parametrize("small_page", [False, True]) @pytest.mark.parametrize("tp_world_size", [1, 4]) @pytest.mark.parametrize("mtp_step", [0, 2]) -def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, tp_world_size, mtp_step): +@pytest.mark.parametrize("kv_type", ["None", "fp8kv_dsa"]) +def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, tp_world_size, mtp_step, kv_type): + from lightllm.models.glm5_next.model import Glm5NextTpPartModel + monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_NODE", "0") monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") monkeypatch.setattr("lightllm.common.req_manager.req_sampling_params.get_vocab_size", lambda _: 128) args = StartArgs( tp=tp_world_size, + llm_kv_type=kv_type, data_type="bfloat16", linear_att_hash_page_size=4, linear_att_page_block_num=2, @@ -54,12 +74,14 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, t set_unique_server_name(args) get_env_start_args.cache_clear() set_env_start_args(dataclasses.asdict(args)) + use_fp8_kv = kv_type == "fp8kv_dsa" + dtype, head_dim = (torch.uint8, 720) if use_fp8_kv else (torch.bfloat16, 584) config = Glm5NextCacheConfig( tp_world_size=tp_world_size, full_att_all_num_kv_heads=1, - full_att_dtype=torch.bfloat16, + full_att_dtype=dtype, full_att_num_kv_heads=1, - full_att_head_dim=584, + full_att_head_dim=head_dim, global_linear_k_heads=2 * tp_world_size, global_linear_v_heads=2 * tp_world_size, num_linear_k_heads=2, @@ -75,18 +97,32 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, t draft_full_att_kv_layer_num=int(mtp_step > 0), ) full_layers = config.get_full_att_kv_layer_num_with_draft_model() - mem = Glm5NextMemManager(16, torch.bfloat16, 1, 584, full_layers, config) + model = object.__new__(Glm5NextTpPartModel) + model.args = args + model.max_total_token_num = 16 + model.linear_config = config + model.mem_fraction = 0.9 + model._init_mem_manager() + mem = model.mem_manager req = Glm5NextReqManager(3, 16, mem, config) att_kv = mem.get_att_input_params(3) - assert att_kv.shape == (17, 1, 512) - assert att_kv.stride(0) == 584 + assert att_kv.shape == (17, 1, 584 if use_fp8_kv else 512) + assert att_kv.stride(0) == head_dim index_bytes = mem.get_indexer_k_buffer(3) index_bytes.random_(0, 256) expected_index_bytes = index_bytes.clone() new_kv = torch.randn(2, 1, 512, dtype=torch.bfloat16, device="cuda") destinations = torch.tensor([5, 9], dtype=torch.int32, device="cuda") mem.operator.copy_kv_to_mem_manager(3, destinations, new_kv) - assert torch.equal(att_kv[destinations], new_kv) + if use_fp8_kv: + values = att_kv[..., :448].view(torch.float8_e4m3fn).float().view(17, 1, 7, 64) + scales = att_kv[..., 576:583].view(torch.float8_e8m0fnu).float().unsqueeze(-1) + tail = att_kv[..., 448:576].view(torch.bfloat16).float() + actual = torch.cat(((values * scales).view(17, 1, 448), tail), dim=-1)[destinations] + torch.testing.assert_close(actual, new_kv.float(), atol=0.13, rtol=0.065) + assert torch.equal(tail[destinations], new_kv[..., 448:].float()) + else: + assert torch.equal(att_kv[destinations], new_kv) assert torch.equal(index_bytes, expected_index_bytes) assert req.get_indexer_tail_buffer(3).shape == (4, 4 + mtp_step, 256) if mtp_step: @@ -133,7 +169,7 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, t packed_bytes[:, 0].random_(0, 256) mem.operator.copy_mem_to_mem(torch.tensor([0]), torch.tensor([7])) assert torch.equal(packed_bytes[:, 0], packed_bytes[:, 7]) - assert mem.get_cell_size() == 584 * 2 * full_layers + assert mem.get_cell_size() == head_dim * dtype.itemsize * full_layers assert config.get_cpu_cache_full_att_bytes() == mem.get_cell_size() * 8 from lightllm.common.basemodel.triton_kernel.linear_att_cpu_cache_copy import ( diff --git a/unit_tests/models/glm5_next/test_fp8_kv.py b/unit_tests/models/glm5_next/test_fp8_kv.py new file mode 100644 index 0000000000..901fde8d69 --- /dev/null +++ b/unit_tests/models/glm5_next/test_fp8_kv.py @@ -0,0 +1,1269 @@ +import weakref +from types import SimpleNamespace + +import pytest +import torch +import triton + +from lightllm.common.basemodel.attention.base_att import AttControl +from lightllm.common.basemodel.infer_struct import InferStateInfo +from lightllm.common.basemodel.attention.nsa.glm5_next import ( + Glm5NextSparsePrefillState, + Glm5NextSparseDecodeState, +) +from lightllm.common.basemodel.attention.nsa.fp8_glm5_next import ( + Fp8Glm5NextSparseAttBackend, + Fp8Glm5NextSparsePrefillState, + Fp8Glm5NextSparseDecodeState, +) +from lightllm.models.glm5_next.triton_kernel.fp8_sparse_decode import ( + _can_use_fp8_sparse_direct_decode, + _fp8_sparse_direct_decode, +) +from lightllm.models.glm5_next.triton_kernel.pad_query_heads import _fused_strided_qpad_kernel +from lightllm.models.glm5_next.triton_kernel.destindex_copy_kv_flashmla_fp8 import destindex_copy_kv_flashmla_fp8 +from lightllm.models.glm5_next.triton_kernel import context_flashattention_nopad_fp8 as fp8_prefill +from lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 import ( + _build_prefill_row_table_kernel, + _gather_prefill_kv_kernel, +) + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def _make_infer_state(**attributes): + state = InferStateInfo() + vars(state).update(attributes) + return state + + +def _make_workspace_backend(batch_max_tokens): + backend = object.__new__(Fp8Glm5NextSparseAttBackend) + backend.model = SimpleNamespace(batch_max_tokens=batch_max_tokens) + backend._prefill_workspaces = [{"smallq": None, "generic": None} for _ in range(2)] + return backend + + +@pytest.mark.parametrize("draft", [False, True]) +@pytest.mark.parametrize("kv_type", ["None", "fp8kv_dsa"]) +def test_main_and_native_mtp_backend_selection(monkeypatch, draft, kv_type): + from lightllm.models.glm5_next.model import Glm5NextTpPartModel + from lightllm.models.glm5_next_mtp.model import Glm5NextMTPModel + + monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") + model = object.__new__(Glm5NextMTPModel if draft else Glm5NextTpPartModel) + model.args = SimpleNamespace(llm_kv_type=kv_type) + model.graph_max_batch_size = 4 + model.max_seq_length = 128 + model._init_att_backend() + infer_state = InferStateInfo() + prefill = model.prefill_att_backend.create_att_prefill_state(infer_state) + decode = model.decode_att_backend.create_att_decode_state(infer_state) + assert type(prefill) is (Fp8Glm5NextSparsePrefillState if kv_type == "fp8kv_dsa" else Glm5NextSparsePrefillState) + assert type(decode) is (Fp8Glm5NextSparseDecodeState if kv_type == "fp8kv_dsa" else Glm5NextSparseDecodeState) + infer_state.prefill_att_state = prefill + infer_state.decode_att_state = decode + infer_state.check_att_states() + assert prefill.infer_state is infer_state + assert decode.infer_state is infer_state + owner_ref = weakref.ref(infer_state) + del infer_state + assert owner_ref() is None + with pytest.raises(ReferenceError): + prefill.infer_state + with pytest.raises(ReferenceError): + decode.infer_state + + +def _dequantize_kv_cache(packed_kv): + values = packed_kv[..., :448].view(torch.float8_e4m3fn).float().reshape(-1, 7, 64) + scales = packed_kv[..., 576:583].view(torch.float8_e8m0fnu).float().reshape(-1, 7, 1) + tail = packed_kv[..., 448:576].view(torch.bfloat16) + return torch.cat(((values * scales).reshape(-1, 1, 448).bfloat16(), tail), dim=-1) + + +def _torch_sparse_attention(q, kv, indices, softmax_scale=0.0625): + outputs = [] + for row in range(q.shape[0]): + keys = kv[indices[row][indices[row] >= 0].long(), 0].float() + outputs.append((q[row].float() @ keys.T * softmax_scale).softmax(-1) @ keys) + return torch.stack(outputs) + + +def _capture_cuda_graph(run): + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + run() + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = run() + return graph, output + + +@pytest.mark.parametrize("index_dtype", [torch.int32, torch.int64]) +@pytest.mark.parametrize("source_d_stride", [1, 2]) +def test_fp8_scatter_quantization_and_gather_prefill(index_dtype, source_d_stride): + torch.manual_seed(53) + # Noncontiguous KV from a fused projection, and a strided packed cache. + raw_source = torch.randn(11, 1, 512 * source_d_stride, device="cuda", dtype=torch.bfloat16) + source = raw_source if source_d_stride == 1 else raw_source[..., ::2] + source[0].zero_() + source[1].mul_(1e-5) + mem_index = torch.randperm(31, device="cuda")[:11].to(index_dtype) + packed_kv = torch.full((31, 1, 720), 71, device="cuda", dtype=torch.uint8) + destindex_copy_kv_flashmla_fp8(source, mem_index, packed_kv) + groups = source[..., :448].float().reshape(11, 1, 7, 64) + scales = torch.exp2(torch.ceil(torch.log2((groups.abs().amax(-1) / 448).clamp_min(1e-4)))) + values = (groups / scales.unsqueeze(-1)).to(torch.float8_e4m3fn) + actual_scales = packed_kv[..., 576:583].view(torch.float8_e8m0fnu).float() + torch.testing.assert_close(actual_scales[mem_index], scales, rtol=0, atol=0) + assert torch.equal(packed_kv[mem_index, :, :448], values.view(torch.uint8).reshape(11, 1, 448)) + assert torch.equal(packed_kv[..., 448:576].view(torch.bfloat16)[mem_index], source[..., 448:]) + assert not packed_kv[mem_index, :, 583].any() + assert (packed_kv[..., 584:] == 71).all() + unused = torch.ones(31, device="cuda", dtype=torch.bool) + unused[mem_index] = False + assert (packed_kv[unused] == 71).all() + + # The packer is graph-safe for both source layouts and mutable destination mapping. + graph, _ = _capture_cuda_graph(lambda: destindex_copy_kv_flashmla_fp8(source, mem_index, packed_kv)) + graph.replay() + assert torch.equal(packed_kv[mem_index, :, :448], values.view(torch.uint8).reshape(11, 1, 448)) + source.add_(0.125) + mem_index.copy_(mem_index.roll(1)) + graph.replay() + changed_groups = source[..., :448].float().reshape(11, 1, 7, 64) + changed_scales = torch.exp2(torch.ceil(torch.log2((changed_groups.abs().amax(-1) / 448).clamp_min(1e-4)))) + changed_values = (changed_groups / changed_scales.unsqueeze(-1)).to(torch.float8_e4m3fn) + assert torch.equal(packed_kv[mem_index, :, :448], changed_values.view(torch.uint8).reshape(11, 1, 448)) + + indices = mem_index[[0, 2, 8, 2, 5, 1, 9, 3]] + prefill_kv = source[8:].clone().add_(0.25) + row_table = torch.full((31,), -1, device="cuda", dtype=torch.int32) + _build_prefill_row_table_kernel[(1,)](mem_index[8:], row_table, 3, BLOCK=256) + gathered_kv = torch.empty((indices.numel(), 1, 512), device="cuda", dtype=torch.bfloat16) + _gather_prefill_kv_kernel[(triton.cdiv(indices.numel(), 4),)]( + packed_kv[..., :584], + indices, + row_table, + prefill_kv, + gathered_kv, + indices.numel(), + packed_kv.stride(0), + prefill_kv.stride(0), + prefill_kv.stride(2), + BLOCK_ROWS=4, + ) + expected = _dequantize_kv_cache(packed_kv) + expected[mem_index[8:]] = prefill_kv + assert torch.equal(gathered_kv, expected[indices.long()]) + + +@pytest.mark.parametrize("num_heads", [8, 16, 64, 128]) +def test_fp8_nope_decode_and_cuda_graph(num_heads): + if torch.cuda.get_device_capability()[0] != 9: + pytest.skip("Validated FlashMLA sparse decode target is Hopper") + import flash_mla + + torch.manual_seed(53) + packed_kv = torch.zeros(4096, 1, 720, device="cuda", dtype=torch.uint8) + source = torch.randn(4096, 1, 512, device="cuda", dtype=torch.bfloat16) + mem_index = torch.arange(4096, device="cuda", dtype=torch.int32) + destindex_copy_kv_flashmla_fp8(source, mem_index, packed_kv) + q = torch.randn(num_heads, 9, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) + # Includes a HOLD row, short tails, and three consecutive MTP positions. + lengths = torch.tensor([0, 1, 3, 127, 128, 511, 2048, 2049, 2051], device="cuda", dtype=torch.int32) + indices = torch.randint(0, 4096, (9, 2176), device="cuda", dtype=torch.int32) + lanes = torch.arange(2176, device="cuda") + indices.masked_fill_(lanes[None, :] >= lengths[:, None], -1) + control = AttControl( + nsa_decode=True, nsa_decode_dict={"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512} + ) + state = Fp8Glm5NextSparseDecodeState(flashmla_sched_meta=flash_mla.get_mla_metadata()[0]) + + def run(): + return state.decode_att((q, q[..., :0]), packed_kv[..., :584], None, control) + + def check(actual): + expected = _torch_sparse_attention(q, _dequantize_kv_cache(packed_kv), indices) + torch.testing.assert_close(actual.float(), expected, atol=0.012, rtol=0.015) + + check(run()) + graph, out = _capture_cuda_graph(run) + graph.replay() + check(out) + q.normal_() + source.normal_() + destindex_copy_kv_flashmla_fp8(source, mem_index, packed_kv) + lengths.copy_(lengths.roll(1)) + indices.random_(0, 4096) + indices.masked_fill_(lanes[None, :] >= lengths[:, None], -1) + graph.replay() + check(out) + + +@pytest.mark.parametrize("cached", [False, True]) +@pytest.mark.parametrize("num_heads", [8, 16, 64, 128]) +@pytest.mark.parametrize("chunk_size", [3, 7, 32768]) +def test_fp8_nope_batched_prefill_and_graph(monkeypatch, cached, num_heads, chunk_size): + monkeypatch.setattr(fp8_prefill, "KV_GATHER_CHUNK_SIZE", chunk_size) + torch.manual_seed(53) + # Two requests whose physical cache locations are unrelated to batch order. + prefix_len = 4 if cached else 0 + seq_len = prefix_len + 3 + ragged_mem_index = torch.randperm(64, device="cuda")[: 2 * seq_len].int() + prefill_rows = torch.cat( + (torch.arange(prefix_len, seq_len), torch.arange(seq_len + prefix_len, 2 * seq_len)) + ).cuda() + prefill_mem_index = ragged_mem_index[prefill_rows] + source = torch.randn(2 * seq_len, 1, 512, device="cuda", dtype=torch.bfloat16) + prefill_kv = source[prefill_rows].clone().add_(0.25) + packed_kv = torch.zeros(64, 1, 720, device="cuda", dtype=torch.uint8) + destindex_copy_kv_flashmla_fp8(source, ragged_mem_index, packed_kv) + q = torch.randn(num_heads, 6, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) + kv_start_offsets = torch.tensor([0] * 3 + [seq_len] * 3, device="cuda", dtype=torch.int32) + topk_indices = torch.full((6, 128), -1, device="cuda", dtype=torch.int32) + for row in range(6): + length = prefix_len + row % 3 + 1 + topk_indices[row, :length] = torch.arange(length, device="cuda") + topk_indices[0].fill_(-1) + topk_indices[1, :2] = prefix_len + topk_indices[3].fill_(-1) + topk_indices[3, 0] = prefix_len + control = AttControl( + nsa_prefill=True, + nsa_prefill_dict={ + "topk_indices": topk_indices, + "prefill_cache_kv": prefill_kv, + "softmax_scale": 512 ** -0.5, + "kv_lora_rank": 512, + }, + ) + infer_state = _make_infer_state(max_cache_len=prefix_len, mem_index=prefill_mem_index) + state = Fp8Glm5NextSparsePrefillState( + _infer_state_weak_ref=weakref.ref(infer_state), + ks=kv_start_offsets, + ragged_mem_index=ragged_mem_index, + ) + + def run(): + return state.prefill_att(q, packed_kv[..., :584], None, control) + + def check(out): + expected_kv = _dequantize_kv_cache(packed_kv)[ragged_mem_index.long()] + expected_kv[prefill_rows] = prefill_kv + indices = torch.where(topk_indices >= 0, topk_indices + kv_start_offsets[:, None], -1) + expected = _torch_sparse_attention(q, expected_kv, indices, control.nsa_prefill_dict["softmax_scale"]) + torch.testing.assert_close(out.float(), expected, atol=0.012, rtol=0.015) + + check(run()) + graph, out = _capture_cuda_graph(run) + graph.replay() + check(out) + q.normal_() + source.normal_() + prefill_kv.copy_(source[prefill_rows] + 0.25) + ragged_mem_index.copy_(ragged_mem_index.roll(1)) + prefill_mem_index.copy_(ragged_mem_index[prefill_rows]) + topk_indices.copy_(topk_indices.roll(1, dims=1)) + destindex_copy_kv_flashmla_fp8(source, ragged_mem_index, packed_kv) + graph.replay() + check(out) + + +@pytest.mark.parametrize("num_heads", [8, 16, 64, 128]) +@pytest.mark.parametrize("softmax_scale", [0.0625, 512 ** -0.5]) +def test_fp8_chunked_prefill_sparse_rows_and_graph(monkeypatch, num_heads, softmax_scale): + monkeypatch.setattr(fp8_prefill, "KV_GATHER_CHUNK_SIZE", 129) + torch.manual_seed(53) + num_tokens = 1025 + source = torch.randn(num_tokens, 1, 1024, device="cuda", dtype=torch.bfloat16)[..., :512] + ragged_mem_index = torch.randperm(2048, device="cuda")[:num_tokens].int() + prefill_rows = torch.tensor([1, 15, 128, 129, 255, 257, 511, 512, 1024], device="cuda") + prefill_mem_index = ragged_mem_index[prefill_rows].long() + prefill_kv = torch.randn(9, 1, 1024, device="cuda", dtype=torch.bfloat16)[..., :512] + packed_kv = torch.zeros(2048, 1, 720, device="cuda", dtype=torch.uint8) + destindex_copy_kv_flashmla_fp8(source, ragged_mem_index, packed_kv) + q = torch.randn(num_heads, 9, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) + lengths = torch.tensor([0, 1, 3, 127, 128, 511, 2048, 2049, 2051], device="cuda") + topk_indices = torch.randint(num_tokens, (9, 2176), device="cuda", dtype=torch.int32) + lanes = torch.arange(2176, device="cuda") + topk_indices.masked_fill_(lanes[None, :] >= lengths[:, None], -1) + topk_indices[1, 0] = 15 + topk_indices[2, :3] = torch.tensor([2, 3, 4], device="cuda") + topk_indices[3, :127] = prefill_rows[torch.randint(9, (127,), device="cuda")].int() + control = AttControl( + nsa_prefill=True, + nsa_prefill_dict={ + "topk_indices": topk_indices, + "prefill_cache_kv": prefill_kv, + "softmax_scale": softmax_scale, + "kv_lora_rank": 512, + }, + ) + infer_state = _make_infer_state(mem_index=prefill_mem_index) + state = Fp8Glm5NextSparsePrefillState( + _infer_state_weak_ref=weakref.ref(infer_state), + ks=torch.zeros(9, device="cuda", dtype=torch.int32), + ragged_mem_index=ragged_mem_index, + ) + + def run(): + return state.prefill_att(q, packed_kv[..., :584], None, control) + + def check(actual): + expected_kv = _dequantize_kv_cache(packed_kv)[ragged_mem_index.long()] + expected_kv[prefill_rows] = prefill_kv + expected = _torch_sparse_attention(q, expected_kv, topk_indices, softmax_scale) + torch.testing.assert_close(actual.float(), expected, atol=0.012, rtol=0.015) + + check(run()) + graph, output = _capture_cuda_graph(run) + graph.replay() + check(output) + q.normal_() + source.normal_() + prefill_kv.normal_() + ragged_mem_index.copy_(ragged_mem_index.roll(1)) + prefill_mem_index.copy_(ragged_mem_index[prefill_rows]) + destindex_copy_kv_flashmla_fp8(source, ragged_mem_index, packed_kv) + topk_indices.copy_(topk_indices.roll(1, dims=1)) + graph.replay() + check(output) + + +def _sampled_h64_sparse_reference(q, packed, ragged, prefill_mem, prefill_kv, starts, topk, samples): + """Reference only a few H64 rows so the large-Q regression stays bounded.""" + expected = [] + for row, heads in samples: + local = topk[row] + local = local[local >= 0].long() + physical = ragged[(starts[row].long() + local)].long() + values = _dequantize_kv_cache(packed[physical])[:, 0].float() + current = physical[:, None] == prefill_mem[None, :].long() + has_current = current.any(dim=1) + current_row = current.int().argmax(dim=1) + values[has_current] = prefill_kv[current_row[has_current], 0].float() + expected.append(((q[row, heads].float() @ values.T * 0.0625).softmax(-1) @ values)) + return expected + + +@pytest.mark.parametrize("num_q", [16, 33]) +def test_fp8_h64_largeq_native_prefill_permuted_ragged_graph(monkeypatch, num_q): + """The native sparse prefill path maps top-k ragged positions, never physical IDs.""" + if torch.cuda.get_device_capability()[0] != 9: + pytest.skip("Validated FlashMLA sparse prefill target is Hopper") + from sgl_kernel.flash_mla import flash_mla_sparse_fwd + + torch.manual_seed(53100 + num_q) + capacity = 131072 + ragged_n = 65536 + packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) + source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) + destindex_copy_kv_flashmla_fp8(source, torch.arange(capacity, device="cuda", dtype=torch.int32), packed) + ragged = torch.randperm(capacity, device="cuda", dtype=torch.int32)[:ragged_n] + q = torch.randn(num_q, 64, 512, device="cuda", dtype=torch.bfloat16) + starts = torch.where( + torch.arange(num_q, device="cuda", dtype=torch.int32) < num_q // 2, + torch.full((num_q,), 1024, device="cuda", dtype=torch.int32), + torch.full((num_q,), 32768, device="cuda", dtype=torch.int32), + ) + current_local = (torch.arange(num_q, device="cuda", dtype=torch.int32) * 37) % 2048 + 64 + current_positions = starts + current_local + prefill_mem = ragged[current_positions.long()].clone() + prefill_kv = torch.randn(num_q, 1, 512, device="cuda", dtype=torch.bfloat16) + prefill_kv.add_(torch.arange(num_q, device="cuda", dtype=torch.bfloat16)[:, None, None]) + topk = torch.randint(0, 2176, (num_q, 2176), device="cuda", dtype=torch.int32) + topk[:, ::19] = -1 # Holes must survive the native position-preserving map. + topk[:, 0] = current_local + topk[:, 1] = current_local # Duplicates are valid selected keys. + topk[:, 2] = topk[:, 3] + assert q.is_contiguous() and topk.is_contiguous() + + observed = [] + + def observed_sparse(*args, **kwargs): + observed.append((kwargs["q"].shape, kwargs["kv"].shape, kwargs["indices"].shape)) + return flash_mla_sparse_fwd(*args, **kwargs) + + monkeypatch.setattr("sgl_kernel.flash_mla.flash_mla_sparse_fwd", observed_sparse) + + def run(): + return fp8_prefill.context_attention_fwd_fp8( + q, packed[..., :584], ragged, prefill_mem, prefill_kv, starts, topk, 0.0625 + ) + + samples = [(0, torch.tensor([0, 9, 63], device="cuda")), (num_q - 1, torch.tensor([3, 31, 62], device="cuda"))] + + def check(actual): + expected = _sampled_h64_sparse_reference(q, packed, ragged, prefill_mem, prefill_kv, starts, topk, samples) + for (row, heads), value in zip(samples, expected): + torch.testing.assert_close(actual[row, heads].float(), value, atol=0.012, rtol=0.015) + + eager = run() + check(eager) + assert observed == [((num_q, 64, 512), (ragged_n, 1, 512), (num_q, 1, 2176))] + graph, output = _capture_cuda_graph(run) + q.normal_() + prefill_kv.normal_() + ragged.copy_(ragged.roll(1)) + prefill_mem.copy_(ragged[current_positions.long()]) + topk.copy_(topk.roll(1, dims=1)) + topk[:, 1] = current_local + graph.replay() + check(output) + topk.fill_(-1) + graph.replay() + assert torch.count_nonzero(output) == 0 + assert len(observed) == 3 # eager, graph warmup, and capture; replay stays inside the CUDA graph. + + +def test_fp8_h8_largeq_prefill_keeps_flashattention_path(monkeypatch): + """H8 large-Q prefill must not enter the H64 native sparse fast path.""" + import sgl_kernel.flash_attn as flash_attn + import sgl_kernel.flash_mla as flash_mla + + torch.manual_seed(53180) + capacity = 65536 + num_q = 16 + packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) + source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) + destindex_copy_kv_flashmla_fp8(source, torch.arange(capacity, device="cuda", dtype=torch.int32), packed) + q = torch.randn(num_q, 8, 512, device="cuda", dtype=torch.bfloat16) + ragged = torch.randperm(capacity, device="cuda", dtype=torch.int32) + starts = torch.full((num_q,), 1024, device="cuda", dtype=torch.int32) + current_local = torch.arange(num_q, device="cuda", dtype=torch.int32) + 64 + prefill_mem = ragged[(starts + current_local).long()].clone() + prefill_kv = torch.randn(num_q, 1, 512, device="cuda", dtype=torch.bfloat16) + topk = torch.full((num_q, 2176), -1, device="cuda", dtype=torch.int32) + topk[:, :64] = torch.arange(64, device="cuda", dtype=torch.int32) + topk[:, 0] = current_local + seen_flashattention = [] + original_flashattention = flash_attn.flash_attn_with_kvcache + + def observed_flashattention(*args, **kwargs): + seen_flashattention.append(True) + return original_flashattention(*args, **kwargs) + + def unexpected_native(*args, **kwargs): + raise AssertionError("H8 large-Q prefill entered the native H64 sparse path") + + monkeypatch.setattr(flash_attn, "flash_attn_with_kvcache", observed_flashattention) + monkeypatch.setattr(flash_mla, "flash_mla_sparse_fwd", unexpected_native) + output = fp8_prefill.context_attention_fwd_fp8( + q, packed[..., :584], ragged, prefill_mem, prefill_kv, starts, topk, 0.0625 + ) + assert seen_flashattention + assert output.shape == q.shape + + +@pytest.mark.parametrize( + "num_q,num_heads", + [(1, 8), (2, 8), (8, 8), (1, 64), (2, 64), (3, 64), (4, 64), (5, 64), (6, 64), (7, 64), (8, 64)], +) +def test_fp8_direct_smallq_prefill_dispatch_strided_holes_and_current(num_q, num_heads): + """Exercise eager and graph replay for no-rowtable, small-Q packed prefill.""" + torch.manual_seed(53016 + num_q) + capacity = 65536 + packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) + source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) + mem = torch.arange(capacity, device="cuda", dtype=torch.int32) + destindex_copy_kv_flashmla_fp8(source, mem, packed) + ragged = torch.arange(capacity, device="cuda", dtype=torch.int32) + q = torch.randn(num_heads, num_q, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) + # Requests deliberately have different spans; starts are not fixed-stride. + seq_lens = torch.tensor([127, 521, 1031, 2053, 4099, 8191, 12289, 16381], device="cuda", dtype=torch.int32)[:num_q] + starts = torch.nn.functional.pad(seq_lens.cumsum(0), (1, 0))[:-1] + current_rows = starts + seq_lens - 1 + prefill_mem = ragged[current_rows] + prefill_kv = torch.randn(num_q, 1, 1024, device="cuda", dtype=torch.bfloat16)[..., :512] + topk_backing = torch.full((num_q, 2176 * 2), -1, device="cuda", dtype=torch.int32) + topk = topk_backing[:, ::2] # Exercise the topk stride passed to the direct kernel. + topk[:, :3] = torch.tensor([3, 11, 11], device="cuda", dtype=torch.int32) + topk[:, 3] = seq_lens - 1 + + def run(): + return fp8_prefill.context_attention_fwd_fp8( + q, packed[..., :584], ragged, prefill_mem, prefill_kv, starts, topk, 0.0625 + ) + + def check(actual): + assert actual.shape == q.shape and actual.is_contiguous() + expected_kv = _dequantize_kv_cache(packed)[ragged.long()] + expected_kv[current_rows.long()] = prefill_kv + physical = torch.where(topk >= 0, topk + starts[:, None], -1) + expected = _torch_sparse_attention(q, expected_kv, physical) + torch.testing.assert_close(actual.float(), expected, atol=0.012, rtol=0.015) + + actual = run() + check(actual) + assert torch.count_nonzero(actual) > 0 # Includes BF16 current KV in every row. + graph, output = _capture_cuda_graph(run) + graph.replay() + check(output) + q.normal_() + source.normal_() + prefill_kv.normal_() + ragged.copy_(ragged.roll(1)) + prefill_mem.copy_(ragged[current_rows]) + destindex_copy_kv_flashmla_fp8(source, mem, packed) + topk.copy_(topk.roll(1, dims=1)) + graph.replay() + check(output) + topk.fill_(-1) # Empty graph replay must not read invalid packed pointers. + graph.replay() + assert torch.count_nonzero(output) == 0 + + +def test_fp8_h64_direct_smallq_state_workspace_reuse_and_fresh_state(): + """Q1 H64 direct workspace is owned by the backend and reused across infer states.""" + torch.manual_seed(53150) + capacity = 65536 + packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) + source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) + mem = torch.arange(capacity, device="cuda", dtype=torch.int32) + destindex_copy_kv_flashmla_fp8(source, mem, packed) + ragged = torch.arange(capacity, device="cuda", dtype=torch.int32) + q = torch.randn(64, 1, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1).contiguous() + starts = torch.zeros(1, device="cuda", dtype=torch.int32) + current_rows = torch.tensor([2048], device="cuda", dtype=torch.int32) + prefill_mem = ragged[current_rows] + prefill_kv = torch.randn(1, 1, 512, device="cuda", dtype=torch.bfloat16) + topk = torch.full((1, 2176), -1, device="cuda", dtype=torch.int32) + topk[:, :2048] = torch.arange(2048, device="cuda", dtype=torch.int32) + topk[:, 2048] = current_rows + control = AttControl( + nsa_prefill=True, + nsa_prefill_dict={ + "topk_indices": topk, + "prefill_cache_kv": prefill_kv, + "softmax_scale": 0.0625, + "kv_lora_rank": 512, + }, + ) + infer_state = _make_infer_state(mem_index=prefill_mem) + backend = _make_workspace_backend(q.shape[0]) + state = Fp8Glm5NextSparsePrefillState( + backend=backend, _infer_state_weak_ref=weakref.ref(infer_state), ks=starts, ragged_mem_index=ragged + ) + + def expected(): + values = _dequantize_kv_cache(packed)[ragged.long()] + values[current_rows.long()] = prefill_kv + return _torch_sparse_attention(q, values, topk) + + def run(active_state=state): + return active_state.prefill_att(q, packed[..., :584], None, control) + + torch.testing.assert_close(run().float(), expected(), atol=0.012, rtol=0.015) + cache = backend._prefill_workspaces[0]["smallq"] + assert not hasattr(infer_state, "_glm5_fp8_smallq_workspace_cache") + workspace = cache["direct"] + pointers = {name: tensor.data_ptr() for name, tensor in workspace.items()} + for _ in range(10): + layer_state = Fp8Glm5NextSparsePrefillState( + backend=backend, _infer_state_weak_ref=weakref.ref(infer_state), ks=starts, ragged_mem_index=ragged + ) + torch.testing.assert_close(run(layer_state).float(), expected(), atol=0.012, rtol=0.015) + assert {name: tensor.data_ptr() for name, tensor in cache["direct"].items()} == pointers + + fresh_infer_state = _make_infer_state(mem_index=prefill_mem.clone()) + fresh_state = Fp8Glm5NextSparsePrefillState( + backend=backend, _infer_state_weak_ref=weakref.ref(fresh_infer_state), ks=starts, ragged_mem_index=ragged + ) + torch.testing.assert_close(run(fresh_state).float(), expected(), atol=0.012, rtol=0.015) + fresh_workspace = backend._prefill_workspaces[0]["smallq"]["direct"] + assert fresh_workspace["partial_o"].data_ptr() == workspace["partial_o"].data_ptr() + + q.normal_() + source.normal_() + prefill_kv.normal_() + destindex_copy_kv_flashmla_fp8(source, mem, packed) + torch.testing.assert_close(run().float(), expected(), atol=0.012, rtol=0.015) + assert {name: tensor.data_ptr() for name, tensor in cache["direct"].items()} == pointers + + +@pytest.mark.parametrize("num_q", [2, 3, 5, 8]) +def test_fp8_h64_selected_prefill_state_workspace_reuse(num_q): + """H64 selected prefill rewrites backend workspace across DSA layers.""" + torch.manual_seed(53200 + num_q) + capacity = 65536 + packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) + source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) + mem = torch.arange(capacity, device="cuda", dtype=torch.int32) + destindex_copy_kv_flashmla_fp8(source, mem, packed) + ragged = torch.arange(capacity, device="cuda", dtype=torch.int32) + q = torch.randn(64, num_q, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) + starts = torch.arange(num_q, device="cuda", dtype=torch.int32) * 4096 + current_rows = starts + 2048 + prefill_mem = ragged[current_rows] + prefill_kv = torch.randn(num_q, 1, 512, device="cuda", dtype=torch.bfloat16) + topk = torch.full((num_q, 2176), -1, device="cuda", dtype=torch.int32) + topk[:, :2048] = torch.arange(2048, device="cuda", dtype=torch.int32) + topk[:, 2048] = 2048 + control = AttControl( + nsa_prefill=True, + nsa_prefill_dict={ + "topk_indices": topk, + "prefill_cache_kv": prefill_kv, + "softmax_scale": 0.0625, + "kv_lora_rank": 512, + }, + ) + infer_state = _make_infer_state(mem_index=prefill_mem) + backend = _make_workspace_backend(q.shape[0]) + state = Fp8Glm5NextSparsePrefillState( + backend=backend, _infer_state_weak_ref=weakref.ref(infer_state), ks=starts, ragged_mem_index=ragged + ) + + def expected(): + values = _dequantize_kv_cache(packed)[ragged.long()] + values[current_rows.long()] = prefill_kv + physical = torch.where(topk >= 0, topk + starts[:, None], -1) + return _torch_sparse_attention(q, values, physical) + + def run(active_state=state): + return active_state.prefill_att(q, packed[..., :584], None, control) + + eager = run() + torch.testing.assert_close(eager.float(), expected(), atol=0.012, rtol=0.015) + cache = backend._prefill_workspaces[0]["smallq"] + assert not hasattr(infer_state, "_glm5_fp8_smallq_workspace_cache") + workspace = cache["selected"] + pointers = {name: tensor.data_ptr() for name, tensor in workspace.items()} + # All 11 DSA layers share the backend allocation and rewrite its contents. + for _ in range(10): + layer_state = Fp8Glm5NextSparsePrefillState( + backend=backend, _infer_state_weak_ref=weakref.ref(infer_state), ks=starts, ragged_mem_index=ragged + ) + torch.testing.assert_close(run(layer_state).float(), expected(), atol=0.012, rtol=0.015) + assert {name: tensor.data_ptr() for name, tensor in cache["selected"].items()} == pointers + + fresh_infer_state = _make_infer_state(mem_index=prefill_mem.clone()) + fresh_state = Fp8Glm5NextSparsePrefillState( + backend=backend, _infer_state_weak_ref=weakref.ref(fresh_infer_state), ks=starts, ragged_mem_index=ragged + ) + torch.testing.assert_close(run(fresh_state).float(), expected(), atol=0.012, rtol=0.015) + fresh_workspace = backend._prefill_workspaces[0]["smallq"]["selected"] + assert fresh_workspace["compact"].data_ptr() == workspace["compact"].data_ptr() + + source.normal_() + destindex_copy_kv_flashmla_fp8(source, mem, packed) + prefill_kv.normal_() + q.normal_() + eager = run() + assert {name: tensor.data_ptr() for name, tensor in cache["selected"].items()} == pointers + torch.testing.assert_close(eager.float(), expected(), atol=0.012, rtol=0.015) + graph, output = _capture_cuda_graph(run) + q.normal_() + topk.copy_(topk.roll(1, dims=1)) + graph.replay() + torch.testing.assert_close(output.float(), expected(), atol=0.012, rtol=0.015) + topk.fill_(-1) + graph.replay() + assert torch.count_nonzero(output) == 0 + + +def test_fp8_h64_b1_decode_uses_native_eager_and_direct_graph(): + if torch.cuda.get_device_capability()[0] != 9: + pytest.skip("Validated FlashMLA sparse decode target is Hopper") + import flash_mla + + torch.manual_seed(53300) + capacity = 4096 + packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) + source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) + mem = torch.arange(capacity, device="cuda", dtype=torch.int32) + destindex_copy_kv_flashmla_fp8(source, mem, packed) + q = torch.randn(64, 1, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) + indices = torch.full((1, 2176), -1, device="cuda", dtype=torch.int32) + indices[:, :2049] = torch.randint(0, capacity, (1, 2049), device="cuda", dtype=torch.int32) + control = AttControl( + nsa_decode=True, nsa_decode_dict={"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512} + ) + infer = _make_infer_state(is_cuda_graph_capture=False) + state = Fp8Glm5NextSparseDecodeState( + _infer_state_weak_ref=weakref.ref(infer), + decode_topk_bound=2051, + flashmla_sched_meta=flash_mla.get_mla_metadata()[0], + ) + + def run(): + return state.decode_att((q, q[..., :0]), packed[..., :584], None, control) + + expected = _torch_sparse_attention(q, _dequantize_kv_cache(packed), indices) + eager = run() + torch.testing.assert_close(eager.float(), expected, atol=0.012, rtol=0.015) + infer.is_cuda_graph_capture = True + graph, output = _capture_cuda_graph(run) + infer.is_cuda_graph_capture = False + graph.replay() + torch.testing.assert_close(output.float(), expected, atol=0.012, rtol=0.015) + + +@pytest.mark.parametrize("num_heads", [8, 16, 64, 128]) +def test_fp8_decode_qpad_kernel_strides_and_zero_padding(num_heads): + """The capture padding kernel reads a true D-stride and writes every padded head as zero.""" + torch.manual_seed(53016 + num_heads) + batch = 2 + # This has BMM-style B/H strides and a non-unit final dimension stride. + q = torch.randn(num_heads, batch, 1024, device="cuda", dtype=torch.bfloat16)[..., ::2].transpose(0, 1) + padded_heads = ((num_heads + 63) // 64) * 64 + head_dim = q.shape[-1] + block_size = triton.next_power_of_2(2 * head_dim) + padded = torch.empty((batch, 1, padded_heads, head_dim), device="cuda", dtype=torch.bfloat16) + _fused_strided_qpad_kernel[(batch, triton.cdiv(padded_heads * head_dim, block_size))]( + q, + padded, + q.stride(0), + q.stride(1), + q.stride(2), + num_heads, + padded_heads, + head_dim=head_dim, + BLOCK_SIZE=block_size, + num_warps=4, + ) + expected = torch.zeros_like(padded) + expected[:, 0, :num_heads] = q + assert torch.equal(padded, expected) + assert torch.count_nonzero(padded[:, 0, num_heads:]) == 0 + + +@pytest.mark.parametrize("num_heads", [8, 16, 64, 128]) +def test_fp8_decode_infer_state_graph_qpad_and_mutation(num_heads): + """Exercise the actual backend's eager and graph flags with strided BMM Q.""" + if torch.cuda.get_device_capability()[0] != 9: + pytest.skip("Validated FlashMLA sparse decode target is Hopper") + import flash_mla + + torch.manual_seed(53100 + num_heads) + batch, capacity = 2, 4096 + source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) + packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) + mem = torch.arange(capacity, device="cuda", dtype=torch.int32) + destindex_copy_kv_flashmla_fp8(source, mem, packed) + q = torch.randn(num_heads, batch, 1024, device="cuda", dtype=torch.bfloat16)[..., ::2].transpose(0, 1) + indices = torch.full((batch, 2176), -1, device="cuda", dtype=torch.int32) + indices[:, :2049] = torch.randint(0, capacity, (batch, 2049), device="cuda", dtype=torch.int32) + control = AttControl( + nsa_decode=True, + nsa_decode_dict={"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512}, + ) + infer_state = _make_infer_state(is_cuda_graph_capture=False) + state = Fp8Glm5NextSparseDecodeState( + _infer_state_weak_ref=weakref.ref(infer_state), flashmla_sched_meta=flash_mla.get_mla_metadata()[0] + ) + + def run(): + return state.decode_att((q, q[..., :0]), packed[..., :584], None, control) + + def reference(): + return _torch_sparse_attention(q, _dequantize_kv_cache(packed), indices) + + eager = run() + torch.testing.assert_close(eager.float(), reference(), atol=0.012, rtol=0.015) + assert eager.shape == q.shape + infer_state.is_cuda_graph_capture = True + graph, output = _capture_cuda_graph(run) + infer_state.is_cuda_graph_capture = False + graph.replay() + torch.testing.assert_close(output.float(), reference(), atol=0.012, rtol=0.015) + q.copy_(q.roll(1, dims=-1)) + saved = indices.clone() + indices[:, :2049].copy_(saved[:, :2049].roll(1, dims=1)) + assert torch.equal(indices[:, 2049:], saved[:, 2049:]) + graph.replay() + torch.testing.assert_close(output.float(), reference(), atol=0.012, rtol=0.015) + + +def test_fp8_sparse_direct_decode_factory_bound_guard(): + backend = object.__new__(Fp8Glm5NextSparseAttBackend) + backend.model = SimpleNamespace(config={"index_topk": 2048, "index_kpool": 4}) + infer_state = InferStateInfo() + assert backend.create_att_decode_state(infer_state).decode_topk_bound == 2051 + backend.model = SimpleNamespace(config={"index_topk": 2048, "index_kpool": 8}) + assert backend.create_att_decode_state(infer_state).decode_topk_bound is None + backend.model = SimpleNamespace(config={"index_topk": 2047, "index_kpool": 4}) + assert backend.create_att_decode_state(infer_state).decode_topk_bound is None + + +@pytest.mark.parametrize("batch", [1, 8]) +def test_fp8_sparse_direct_decode_actual_entry_edges_and_graph(monkeypatch, batch): + if torch.cuda.get_device_capability()[0] != 9: + pytest.skip("Validated FlashMLA sparse decode target is Hopper") + import flash_mla + + torch.manual_seed(53200 + batch) + capacity = 4096 + source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) + packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) + mem = torch.arange(capacity, device="cuda", dtype=torch.int32) + destindex_copy_kv_flashmla_fp8(source, mem, packed) + q = torch.randn(8, batch, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) + indices = torch.full((batch, 2176), -1, device="cuda", dtype=torch.int32) + counts = [2051] if batch == 1 else [0, 2048, 2049, 2050, 2051, 1985, 2051, 2051] + for row in range(batch): + count = counts[row] + if count: + indices[row, :count] = torch.randint(0, capacity, (count,), device="cuda", dtype=torch.int32) + if batch == 8: + indices[5, :64] = -1 # holes + indices[6, 64:128] = indices[6, :64] # duplicates + assert torch.equal(indices[:, 2051:], torch.full_like(indices[:, 2051:], -1)) + control = AttControl( + nsa_decode=True, nsa_decode_dict={"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512} + ) + state_owner = _make_infer_state(is_cuda_graph_capture=False) + state = Fp8Glm5NextSparseDecodeState( + _infer_state_weak_ref=weakref.ref(state_owner), + flashmla_sched_meta=flash_mla.get_mla_metadata()[0], + decode_topk_bound=2051, + ) + + def native_must_not_run(*args, **kwargs): + raise AssertionError("supported GLM-5.3 decode took native FlashMLA fallback") + + def run(): + return state.decode_att((q, q[..., :0]), packed[..., :584], None, control) + + assert not _can_use_fp8_sparse_direct_decode(q, packed, indices.unsqueeze(1), 0.0625, 512, 2051) + expected = _torch_sparse_attention(q, _dequantize_kv_cache(packed), indices) + eager = run() + torch.testing.assert_close(eager.float(), expected, atol=0.02, rtol=0.02) + if batch == 1: + assert torch.count_nonzero(eager) > 0 + monkeypatch.setattr(flash_mla, "flash_mla_with_kvcache", native_must_not_run) + state.infer_state.is_cuda_graph_capture = True + graph, output = _capture_cuda_graph(run) + state.infer_state.is_cuda_graph_capture = False + graph.replay() + torch.testing.assert_close(output.float(), expected, atol=0.02, rtol=0.02) + q.copy_(q.roll(1, dims=-1)) + saved = indices.clone() + nonempty = (indices >= 0).sum(dim=1) + for row, count in enumerate(nonempty.tolist()): + if count: + indices[row, :count].copy_(saved[row, :count].roll(1)) + assert torch.equal(indices[:, 2051:], saved[:, 2051:]) + graph.replay() + expected = _torch_sparse_attention(q, _dequantize_kv_cache(packed), indices) + torch.testing.assert_close(output.float(), expected, atol=0.02, rtol=0.02) + indices.fill_(-1) + graph.replay() + assert torch.count_nonzero(output) == 0 + + +@pytest.mark.parametrize( + "num_heads,batch", + [ + (1, 1), + (4, 3), + (7, 5), + (8, 2), + (8, 3), + (8, 4), + (8, 5), + (16, 1), + (16, 3), + (16, 4), + (17, 3), + (32, 1), + (32, 2), + (32, 3), + (64, 1), + (64, 2), + (128, 1), + ], +) +def test_fp8_sparse_direct_decode_general_graph_factory_entry(monkeypatch, num_heads, batch): + """General head/token shapes use the direct path selected by the central predicate.""" + if torch.cuda.get_device_capability()[0] != 9: + pytest.skip("Validated FlashMLA sparse decode target is Hopper") + import flash_mla + + torch.manual_seed(53300 + batch) + capacity = 4096 + source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) + packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) + mem = torch.arange(capacity, device="cuda", dtype=torch.int32) + destindex_copy_kv_flashmla_fp8(source, mem, packed) + # BMM-style B/H stride: (512, batch * 512, 1). + q = torch.randn(num_heads, batch, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1) + assert q.stride() == (512, batch * 512, 1) + indices = torch.full((batch, 2176), -1, device="cuda", dtype=torch.int32) + counts = [([2051, 0, 2048, 2050, 1, 3, 127, 128][row % 8]) for row in range(batch)] + for row, count in enumerate(counts): + if count: + indices[row, :count] = torch.randint(0, capacity, (count,), device="cuda", dtype=torch.int32) + if batch >= 3: + indices[2, :32] = -1 # holes + if batch >= 4: + indices[3, 64:128] = indices[3, :64] # duplicates + assert torch.equal(indices[:, 2051:], torch.full_like(indices[:, 2051:], -1)) + control = AttControl( + nsa_decode=True, + nsa_decode_dict={"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512}, + ) + backend = object.__new__(Fp8Glm5NextSparseAttBackend) + backend.model = SimpleNamespace(config={"index_topk": 2048, "index_kpool": 4}) + infer_state = _make_infer_state(is_cuda_graph_capture=False) + state = backend.create_att_decode_state(infer_state) + state.flashmla_sched_meta = flash_mla.get_mla_metadata()[0] + assert state.decode_topk_bound == 2051 + native_state_owner = _make_infer_state(is_cuda_graph_capture=False) + native_state = Fp8Glm5NextSparseDecodeState( + _infer_state_weak_ref=weakref.ref(native_state_owner), + flashmla_sched_meta=flash_mla.get_mla_metadata()[0], + decode_topk_bound=None, + ) + + def run(current_state): + return current_state.decode_att((q, q[..., :0]), packed[..., :584], None, control) + + native_calls = [] + original_native = flash_mla.flash_mla_with_kvcache + + def observe_native(*args, **kwargs): + native_calls.append(True) + return original_native(*args, **kwargs) + + monkeypatch.setattr(flash_mla, "flash_mla_with_kvcache", observe_native) + eager = run(state) + native_output = run(native_state) + assert native_calls + assert torch.equal(eager, native_output) + expected = _torch_sparse_attention(q, _dequantize_kv_cache(packed), indices) + torch.testing.assert_close(eager.float(), expected, atol=0.02, rtol=0.02) + + def native_must_not_run(*args, **kwargs): + raise AssertionError("supported graph shape did not use the direct packed-KV path") + + monkeypatch.setattr(flash_mla, "flash_mla_with_kvcache", native_must_not_run) + state.infer_state.is_cuda_graph_capture = True + graph, output = _capture_cuda_graph(lambda: run(state)) + state.infer_state.is_cuda_graph_capture = False + graph.replay() + torch.testing.assert_close(output.float(), expected, atol=0.02, rtol=0.02) + q.copy_(q.roll(1, dims=-1)) + saved = indices.clone() + for row, count in enumerate((indices >= 0).sum(dim=1).tolist()): + if count: + indices[row, :count].copy_(saved[row, :count].roll(1)) + assert torch.equal(indices[:, 2051:], saved[:, 2051:]) + graph.replay() + torch.testing.assert_close( + output.float(), _torch_sparse_attention(q, _dequantize_kv_cache(packed), indices), atol=0.02, rtol=0.02 + ) + indices.fill_(-1) + graph.replay() + assert torch.count_nonzero(output) == 0 + + +def test_fp8_sparse_direct_decode_unknown_bound_fallback_and_strided_q(monkeypatch): + if torch.cuda.get_device_capability()[0] != 9: + pytest.skip("Validated FlashMLA sparse decode target is Hopper") + import flash_mla + + torch.manual_seed(53250) + capacity = 4096 + source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) + packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) + mem = torch.arange(capacity, device="cuda", dtype=torch.int32) + destindex_copy_kv_flashmla_fp8(source, mem, packed) + q = torch.randn(1, 8, 512, device="cuda", dtype=torch.bfloat16) + late = torch.full((1, 2176), -1, device="cuda", dtype=torch.int32) + late[:, 2175] = 17 + control = AttControl( + nsa_decode=True, nsa_decode_dict={"topk_mem_indices": late, "softmax_scale": 0.0625, "kv_lora_rank": 512} + ) + calls = [] + original = flash_mla.flash_mla_with_kvcache + + def native(*args, **kwargs): + calls.append(True) + return original(*args, **kwargs) + + monkeypatch.setattr(flash_mla, "flash_mla_with_kvcache", native) + unknown_owner = _make_infer_state(is_cuda_graph_capture=False) + unknown = Fp8Glm5NextSparseDecodeState( + _infer_state_weak_ref=weakref.ref(unknown_owner), flashmla_sched_meta=flash_mla.get_mla_metadata()[0] + ) + actual = unknown.decode_att((q, q[..., :0]), packed[..., :584], None, control) + torch.testing.assert_close( + actual.float(), _torch_sparse_attention(q, _dequantize_kv_cache(packed), late), atol=0.02, rtol=0.02 + ) + assert calls and not _can_use_fp8_sparse_direct_decode( + q, packed[..., :584], late.unsqueeze(1), 0.0625, 512, None, is_cuda_graph_capture=True + ) + q_d2 = torch.randn(1, 8, 1024, device="cuda", dtype=torch.bfloat16)[..., ::2] + good = torch.full((1, 2176), -1, device="cuda", dtype=torch.int32) + good[:, :2051] = torch.randint(0, capacity, (1, 2051), device="cuda", dtype=torch.int32) + assert _can_use_fp8_sparse_direct_decode( + q_d2, packed[..., :584], good.unsqueeze(1), 0.0625, 512, 2051, is_cuda_graph_capture=True + ) + assert not _can_use_fp8_sparse_direct_decode( + q, packed[..., :584], good.unsqueeze(1), 0.0625, 256, 2051, is_cuda_graph_capture=True + ) + + +@pytest.mark.parametrize("num_heads,batch", [(1, 3), (7, 5), (16, 16), (32, 7), (65, 3), (128, 8)]) +@pytest.mark.parametrize("index_dtype", [torch.int32, torch.int64]) +def test_fp8_sparse_direct_decode_strided_inputs_and_packed_row_padding(monkeypatch, num_heads, batch, index_dtype): + """The kernel handles partial head tiles, larger batches and actual Q/index strides.""" + torch.manual_seed(53600 + num_heads + batch) + capacity = 4096 + source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) + packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) + destindex_copy_kv_flashmla_fp8(source, torch.arange(capacity, device="cuda", dtype=torch.int32), packed) + q = torch.randn(num_heads, batch, 1024, device="cuda", dtype=torch.bfloat16)[..., ::2].transpose(0, 1) + indices = torch.full((batch, 1, 4352), -1, device="cuda", dtype=index_dtype)[..., ::2] + indices[..., :2051] = torch.randint(0, capacity, (batch, 1, 2051), device="cuda", dtype=index_dtype) + indices[0].fill_(-1) + + def run(): + return _fp8_sparse_direct_decode(q, packed, indices, 0.0625) + + def expected(): + return _torch_sparse_attention(q, _dequantize_kv_cache(packed), indices[:, 0]) + + if (num_heads, batch) in ((1, 3), (7, 5)): + infer = _make_infer_state(is_cuda_graph_capture=True) + state = Fp8Glm5NextSparseDecodeState(_infer_state_weak_ref=weakref.ref(infer), decode_topk_bound=2051) + + def native_must_not_run(*args, **kwargs): + raise AssertionError("supported strided decode took native fallback") + + monkeypatch.setattr("flash_mla.flash_mla_with_kvcache", native_must_not_run) + torch.testing.assert_close( + state.decode_att( + (q, q[..., :0]), + packed, + None, + AttControl( + nsa_decode=True, + nsa_decode_dict={"topk_mem_indices": indices, "softmax_scale": 0.0625, "kv_lora_rank": 512}, + ), + ).float(), + expected(), + atol=0.02, + rtol=0.02, + ) + + torch.testing.assert_close(run().float(), expected(), atol=0.02, rtol=0.02) + graph, output = _capture_cuda_graph(run) + graph.replay() + torch.testing.assert_close(output.float(), expected(), atol=0.02, rtol=0.02) + q.normal_() + indices[..., :2051].copy_(indices[..., :2051].roll(1, dims=0)) + graph.replay() + torch.testing.assert_close(output.float(), expected(), atol=0.02, rtol=0.02) + + +@pytest.mark.parametrize("index_dtype", [torch.int32, torch.int64]) +def test_fp8_sparse_direct_decode_central_predicate_accepts_strides_and_rejects_unaligned_kv(index_dtype): + q = torch.empty(1, 8, 1024, device="cuda", dtype=torch.bfloat16)[..., ::2] + packed = torch.empty(8, 1, 720, device="cuda", dtype=torch.uint8) + topk = torch.empty(1, 1, 4352, device="cuda", dtype=index_dtype)[..., ::2] + assert _can_use_fp8_sparse_direct_decode(q, packed, topk, 0.0625, 512, 2051, is_cuda_graph_capture=True) + odd_rows = torch.empty(8, 1, 721, device="cuda", dtype=torch.uint8)[..., :584] + odd_base = torch.empty(8, 1, 722, device="cuda", dtype=torch.uint8)[..., 1:585] + strided_columns = torch.empty(8, 1, 1440, device="cuda", dtype=torch.uint8)[..., ::2] + for invalid in (odd_rows, odd_base, strided_columns): + assert not _can_use_fp8_sparse_direct_decode(q, invalid, topk, 0.0625, 512, 2051, is_cuda_graph_capture=True) + for scale in (0.0, -1.0, float("inf"), float("nan")): + assert not _can_use_fp8_sparse_direct_decode(q, packed, topk, scale, 512, 2051, is_cuda_graph_capture=True) + + +@pytest.mark.parametrize("strided_input", ["q", "topk"]) +def test_fp8_sparse_direct_decode_offsets_beyond_int32(strided_input): + batch, heads, capacity = 4, 1, 64 + stride = 800_000_000 + shape = (batch, heads, 512) if strided_input == "q" else (batch, 1, 2176) + dtype = torch.bfloat16 if strided_input == "q" else torch.int32 + storage_size = (batch - 1) * stride + shape[-1] + required_bytes = storage_size * torch.empty((), dtype=dtype).element_size() + torch.cuda.empty_cache() + if torch.cuda.mem_get_info()[0] < required_bytes + 512 * 1024 ** 2: + pytest.skip("Large-stride regression requires sufficient free GPU memory") + # Allocate address space without initializing the gaps between logical rows. + storage = torch.empty(storage_size, device="cuda", dtype=dtype) + strided = storage.as_strided(shape, (stride, shape[-1], 1)) + q = strided if strided_input == "q" else torch.empty(batch, heads, 512, device="cuda", dtype=torch.bfloat16) + topk = strided if strided_input == "topk" else torch.empty(batch, 1, 2176, device="cuda", dtype=torch.int32) + q.copy_(torch.randn(q.shape, device="cuda", dtype=q.dtype)) + topk.fill_(-1) + topk[..., :7] = torch.randint(0, capacity, (batch, 1, 7), device="cuda", dtype=topk.dtype) + packed = torch.empty(capacity, 1, 720, device="cuda", dtype=torch.uint8) + source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) + destindex_copy_kv_flashmla_fp8(source, torch.arange(capacity, device="cuda", dtype=torch.int32), packed) + assert _can_use_fp8_sparse_direct_decode(q, packed, topk, 0.0625, 512, 2051, is_cuda_graph_capture=True) + actual = _fp8_sparse_direct_decode(q, packed, topk, 0.0625) + expected = _torch_sparse_attention(q, _dequantize_kv_cache(packed), topk[:, 0]) + torch.testing.assert_close(actual.float(), expected, atol=0.02, rtol=0.02) + + +@pytest.mark.parametrize("num_heads,batch", [(16, 16), (32, 8), (64, 4), (128, 2)]) +def test_fp8_sparse_direct_decode_large_workload_retains_native_graph(monkeypatch, num_heads, batch): + if torch.cuda.get_device_capability()[0] != 9: + pytest.skip("Validated FlashMLA sparse decode target is Hopper") + import flash_mla + + torch.manual_seed(53700 + num_heads + batch) + packed = torch.empty(4096, 1, 720, device="cuda", dtype=torch.uint8) + source = torch.randn(4096, 1, 512, device="cuda", dtype=torch.bfloat16) + destindex_copy_kv_flashmla_fp8(source, torch.arange(4096, device="cuda", dtype=torch.int32), packed) + q = torch.randn(batch, num_heads, 512, device="cuda", dtype=torch.bfloat16) + topk = torch.full((batch, 1, 2176), -1, device="cuda", dtype=torch.int32) + topk[..., :2051] = torch.randint(0, 4096, (batch, 1, 2051), device="cuda", dtype=torch.int32) + assert not _can_use_fp8_sparse_direct_decode(q, packed, topk, 0.0625, 512, 2051, is_cuda_graph_capture=True) + + def direct_must_not_run(*args, **kwargs): + raise AssertionError("larger workload took the slower direct path") + + monkeypatch.setattr( + "lightllm.models.glm5_next.triton_kernel.fp8_sparse_decode._fp8_sparse_direct_decode", direct_must_not_run + ) + infer = _make_infer_state(is_cuda_graph_capture=True) + state = Fp8Glm5NextSparseDecodeState( + _infer_state_weak_ref=weakref.ref(infer), + decode_topk_bound=2051, + flashmla_sched_meta=flash_mla.get_mla_metadata()[0], + ) + control = AttControl( + nsa_decode=True, nsa_decode_dict={"topk_mem_indices": topk, "softmax_scale": 0.0625, "kv_lora_rank": 512} + ) + run = lambda: state.decode_att((q, q[..., :0]), packed[..., :584], None, control) + expected = _torch_sparse_attention(q, _dequantize_kv_cache(packed), topk[:, 0]) + torch.testing.assert_close(run().float(), expected, atol=0.02, rtol=0.02) + graph, output = _capture_cuda_graph(run) + graph.replay() + torch.testing.assert_close(output.float(), expected, atol=0.02, rtol=0.02) + + +def test_fp8_h64_generic_prefill_workspace_reuse_mutable_inputs_and_graph_fallback(): + """Generic H64 backend buffers are reused, batch metadata resets, and graph falls back.""" + torch.manual_seed(53450) + capacity, num_q = 65536, 13 + packed = torch.empty((capacity, 1, 720), device="cuda", dtype=torch.uint8) + source = torch.randn(capacity, 1, 512, device="cuda", dtype=torch.bfloat16) + mem = torch.arange(capacity, device="cuda", dtype=torch.int32) + destindex_copy_kv_flashmla_fp8(source, mem, packed) + ragged = torch.randperm(capacity, device="cuda", dtype=torch.int32) + starts = torch.arange(num_q, device="cuda", dtype=torch.int32) * 2048 + current_local = torch.full((num_q,), 2048, device="cuda", dtype=torch.int32) + prefill_mem = ragged[(starts + current_local).long()].clone() + prefill_kv = torch.randn(num_q, 1, 512, device="cuda", dtype=torch.bfloat16).add_(2) + q = torch.randn(64, num_q, 512, device="cuda", dtype=torch.bfloat16).transpose(0, 1).contiguous() + topk = torch.full((num_q, 2176), -1, device="cuda", dtype=torch.int32) + topk[:, :2048] = torch.arange(2048, device="cuda", dtype=torch.int32) + topk[:, 2048] = current_local + topk[:, 2049] = current_local + control = AttControl( + nsa_prefill=True, + nsa_prefill_dict={ + "topk_indices": topk, + "prefill_cache_kv": prefill_kv, + "softmax_scale": 0.0625, + "kv_lora_rank": 512, + }, + ) + infer = _make_infer_state(mem_index=prefill_mem) + backend = _make_workspace_backend(2 * num_q) + + def expected(active_ragged=ragged): + values = _dequantize_kv_cache(packed)[active_ragged.long()] + values[(starts + current_local).long()] = prefill_kv + physical = torch.where(topk >= 0, topk + starts[:, None], -1) + return _torch_sparse_attention(q, values, physical) + + def run(active): + return active.prefill_att(q, packed[..., :584], None, control) + + state = Fp8Glm5NextSparsePrefillState( + backend=backend, _infer_state_weak_ref=weakref.ref(infer), ks=starts, ragged_mem_index=ragged + ) + torch.testing.assert_close(run(state).float(), expected(), atol=0.012, rtol=0.015) + cache = backend._prefill_workspaces[0] + assert not hasattr(infer, "_fp8_sparse_generic_prefill_workspace_cache") + workspace = cache["generic"]["workspace"] + pointers = {k: v.data_ptr() for k, v in workspace.items() if isinstance(v, torch.Tensor)} + for _ in range(10): + layer = Fp8Glm5NextSparsePrefillState( + backend=backend, _infer_state_weak_ref=weakref.ref(infer), ks=starts, ragged_mem_index=ragged + ) + torch.testing.assert_close(run(layer).float(), expected(), atol=0.012, rtol=0.015) + assert workspace["row_table_built"] + assert { + k: v.data_ptr() for k, v in cache["generic"]["workspace"].items() if isinstance(v, torch.Tensor) + } == pointers + + # Change raw KV/top-k in place: reusable metadata must not cache either. + q.normal_() + prefill_kv.normal_() + topk.copy_(topk.roll(1, dims=1)) + topk[:, 2049] = current_local + torch.testing.assert_close(run(state).float(), expected(), atol=0.012, rtol=0.015) + + fresh_ragged = ragged.roll(1) + fresh_mem = fresh_ragged[(starts + current_local).long()].clone() + fresh = _make_infer_state(mem_index=fresh_mem) + fresh_state = Fp8Glm5NextSparsePrefillState( + backend=backend, _infer_state_weak_ref=weakref.ref(fresh), ks=starts, ragged_mem_index=fresh_ragged + ) + torch.testing.assert_close(run(fresh_state).float(), expected(fresh_ragged), atol=0.012, rtol=0.015) + assert cache["generic"]["workspace"]["prefill_row_table"].data_ptr() == workspace["prefill_row_table"].data_ptr() + + graph_infer = _make_infer_state(mem_index=prefill_mem, prefill_cuda_graph_exe_list=[]) + graph_state = Fp8Glm5NextSparsePrefillState( + backend=backend, _infer_state_weak_ref=weakref.ref(graph_infer), ks=starts, ragged_mem_index=ragged + ) + torch.testing.assert_close(run(graph_state).float(), expected(), atol=0.012, rtol=0.015) + assert not hasattr(graph_infer, "_fp8_sparse_generic_prefill_workspace_cache") + + explicit_graph_infer = _make_infer_state(mem_index=prefill_mem, is_cuda_graph_capture=True) + explicit_graph_state = Fp8Glm5NextSparsePrefillState( + backend=backend, _infer_state_weak_ref=weakref.ref(explicit_graph_infer), ks=starts, ragged_mem_index=ragged + ) + torch.testing.assert_close(run(explicit_graph_state).float(), expected(), atol=0.012, rtol=0.015) + assert not hasattr(explicit_graph_infer, "_fp8_sparse_generic_prefill_workspace_cache") + + +def test_fp8_generic_prefill_workspace_graph_query_is_generic_h64_only(monkeypatch): + """H8/direct/selected prefill never queries CUDA graph-capture state.""" + import lightllm.models.glm5_next.triton_kernel.context_flashattention_nopad_fp8 as fp8_prefill + + queries = [] + + def query_spy(): + queries.append(True) + raise AssertionError("non-generic path queried CUDA graph capture state") + + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", query_spy) + kv = torch.empty((1, 1, 584), dtype=torch.uint8) + params = { + "prefill_cache_kv": torch.empty((1, 1, 512), dtype=torch.bfloat16), + "topk_indices": torch.empty((1, 2176), dtype=torch.int32), + "softmax_scale": 0.0625, + } + control = AttControl(nsa_prefill=True, nsa_prefill_dict=params) + + def fake_context(*args): + return args[0] + + monkeypatch.setattr(fp8_prefill, "context_attention_fwd_fp8", fake_context) + for heads, direct, selected in ((8, False, False), (64, True, False), (64, False, True)): + monkeypatch.setattr(fp8_prefill, "_can_use_direct_smallq", lambda *args, d=direct: d) + monkeypatch.setattr(fp8_prefill, "_can_use_selected_prefill", lambda *args, s=selected: s) + infer = _make_infer_state(mem_index=torch.empty(1, dtype=torch.int32)) + state = Fp8Glm5NextSparsePrefillState( + _infer_state_weak_ref=weakref.ref(infer), + ks=torch.zeros(1, dtype=torch.int32), + ragged_mem_index=torch.zeros(1, dtype=torch.int32), + ) + state._smallq_workspace = lambda active_q, kind="direct" if direct else "selected": {"kind": kind} + assert state._nsa_prefill_att(torch.empty((1, heads, 512), dtype=torch.bfloat16), kv, control).shape[1] == heads + assert not queries diff --git a/unit_tests/models/glm5_next/test_layers.py b/unit_tests/models/glm5_next/test_layers.py index 7b60c9a267..08f8579174 100644 --- a/unit_tests/models/glm5_next/test_layers.py +++ b/unit_tests/models/glm5_next/test_layers.py @@ -345,6 +345,24 @@ def test_glm_weight_loading_requires_hf(config, model_class): model._verify_params() +@pytest.mark.parametrize( + "kv_type, data_type, kv_lora_rank", + [ + ("fp8kv_sph", torch.bfloat16, 512), + ("fp8kv_spt", torch.bfloat16, 512), + ("fp8kv_dsa", torch.float16, 512), + ("fp8kv_dsa", torch.bfloat16, 256), + ], +) +def test_glm_rejects_unsupported_kv_config(config, kv_type, data_type, kv_lora_rank): + model = object.__new__(Glm5NextTpPartModel) + model.config = dict(config, kv_lora_rank=kv_lora_rank) + model.tp_world_size_, model.data_type = 1, data_type + model.args, model.load_way = StartArgs(llm_kv_type=kv_type), "HF" + with pytest.raises(ValueError, match="GLM-5.3 Flash"): + model._verify_params() + + @pytest.mark.parametrize("tp_world_size", [1, 2]) def test_glm_rejects_tpsp(config, tp_world_size): model = object.__new__(Glm5NextTpPartModel) diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index 57665cbbe3..8817df03ef 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -3,7 +3,7 @@ import pytest import torch -from lightllm.common.kv_cache_mem_manager import Glm5NextMemManager, MemoryManager +from lightllm.common.kv_cache_mem_manager import Glm5NextMemManager, FP8Glm5NextMemManager, MemoryManager from lightllm.common.req_manager import Glm5NextReqManager from lightllm.common.state_cache_manager import Glm5NextCacheConfig from lightllm.models.glm5_next.triton_kernel.kpool import compress_pools @@ -14,14 +14,17 @@ pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.fixture(params=[0, 2]) +@pytest.fixture(params=[(0, False), (2, False), (0, True), (2, True)]) def make_mems(monkeypatch, request): monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_NODE", "0") monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") monkeypatch.setattr("lightllm.common.req_manager.req_sampling_params.get_vocab_size", lambda _: 128) - mtp_step = request.param + mtp_step, use_fp8_kv = request.param + dtype, head_dim = (torch.uint8, 720) if use_fp8_kv else (torch.bfloat16, 584) + mem_manager_class = FP8Glm5NextMemManager if use_fp8_kv else Glm5NextMemManager args = StartArgs( data_type="bfloat16", + llm_kv_type="fp8kv_dsa" if use_fp8_kv else "None", linear_att_hash_page_size=4, linear_att_page_block_num=2, mtp_step=mtp_step, @@ -43,9 +46,9 @@ def create(tp): config = Glm5NextCacheConfig( tp_world_size=tp, full_att_all_num_kv_heads=1, - full_att_dtype=torch.bfloat16, + full_att_dtype=dtype, full_att_num_kv_heads=1, - full_att_head_dim=584, + full_att_head_dim=head_dim, global_linear_k_heads=8, global_linear_v_heads=8, num_linear_k_heads=8 // tp, @@ -62,12 +65,12 @@ def create(tp): ) mems = [] for _ in range(tp): - mem = Glm5NextMemManager(32, torch.bfloat16, 1, 584, 2 + int(mtp_step > 0), config) + mem = mem_manager_class(32, dtype, 1, head_dim, 2 + int(mtp_step > 0), config) req = Glm5NextReqManager(3, 32, mem, config) mem.write_to_shm(req) assert mem.big_page_buffers is not None mems.append(mem) - mems[0].alloc_paged_kv_move_buffer(1, 2048) + mems[0].alloc_paged_kv_move_buffer(1, 4096) return mems return create @@ -83,7 +86,8 @@ def test_pd_kv_and_runtime_state_roundtrip(make_mems, prefill_tp, decode_tp, rem length = 8 + remainder src_indexes = torch.randperm(32)[:length].tolist() dst_indexes = torch.randperm(32)[:length].tolist() - packed = torch.randint(0, 256, (layers, length, 1, 1168), dtype=torch.uint8, device="cuda") + token_bytes = source[0].head_dim * source[0].dtype.itemsize + packed = torch.randint(0, 256, (layers, length, 1, token_bytes), dtype=torch.uint8, device="cuda") conv = torch.randn(6, 3, 8, 128, 3, dtype=torch.bfloat16, device="cuda") ssm = torch.randn(6, 8, 128, 128, dtype=torch.float32, device="cuda") # Non-live slots deliberately contain stale data. The sequence length, not @@ -170,4 +174,4 @@ def test_pd_page_capacity_includes_tail_without_resizing(make_mems): with pytest.raises(AssertionError, match="smaller than global linear att state"): helper.assert_page_size() shape = mem.get_paged_kv_move_buffer_shape(2, 2048) - assert shape == (2, 2048, mem.layer_num, 1, 584) + assert shape == (2, 2048, mem.layer_num, 1, mem.head_dim)