From e7a681d24829475402215c4ad1494da925a6680e Mon Sep 17 00:00:00 2001 From: Codex Date: Tue, 4 Aug 2026 21:42:16 +0800 Subject: [PATCH 01/48] feat: vendor lightllm mla kernels --- .../triton_kernel/mla/LICENSE.lightllm | 201 ++++++ src/sparsevllm/triton_kernel/mla/README.md | 40 ++ src/sparsevllm/triton_kernel/mla/__init__.py | 40 ++ .../triton_kernel/mla/copy_latent.py | 219 +++++++ .../triton_kernel/mla/decode_schedule.py | 398 ++++++++++++ .../triton_kernel/mla/decode_stage1.py | 391 ++++++++++++ .../triton_kernel/mla/decode_stage2.py | 175 ++++++ .../triton_kernel/mla/gather_latent.py | 336 ++++++++++ tests/test_mla_kernels.py | 580 ++++++++++++++++++ 9 files changed, 2380 insertions(+) create mode 100644 src/sparsevllm/triton_kernel/mla/LICENSE.lightllm create mode 100644 src/sparsevllm/triton_kernel/mla/README.md create mode 100644 src/sparsevllm/triton_kernel/mla/__init__.py create mode 100644 src/sparsevllm/triton_kernel/mla/copy_latent.py create mode 100644 src/sparsevllm/triton_kernel/mla/decode_schedule.py create mode 100644 src/sparsevllm/triton_kernel/mla/decode_stage1.py create mode 100644 src/sparsevllm/triton_kernel/mla/decode_stage2.py create mode 100644 src/sparsevllm/triton_kernel/mla/gather_latent.py create mode 100644 tests/test_mla_kernels.py diff --git a/src/sparsevllm/triton_kernel/mla/LICENSE.lightllm b/src/sparsevllm/triton_kernel/mla/LICENSE.lightllm new file mode 100644 index 00000000..261eeb9e --- /dev/null +++ b/src/sparsevllm/triton_kernel/mla/LICENSE.lightllm @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/src/sparsevllm/triton_kernel/mla/README.md b/src/sparsevllm/triton_kernel/mla/README.md new file mode 100644 index 00000000..c931206f --- /dev/null +++ b/src/sparsevllm/triton_kernel/mla/README.md @@ -0,0 +1,40 @@ +# MLA Triton kernels + +This directory contains the minimal LightLLM-derived mathematical kernels used +by Sparse-vLLM's GLM-4.7 MLA path. It deliberately has no dependency on the +LightLLM Python package or runtime. + +## Upstream source + +- Repository: +- Commit: `65c174ee95ac6a6fd36b18b63d0b33d97e76b770` +- License: Apache-2.0; see `LICENSE.lightllm` + +| Local file | Upstream file | Local changes | +| --- | --- | --- | +| `decode_stage1.py` | `lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding_stage1.py` | Removed device probing and runtime state; caller supplies strides, workspace, and static launch values; restricted shapes to GLM's 512+64 latent layout. | +| `decode_stage2.py` | `lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding_stage2.py` | Removed runtime imports; caller supplies schedule/workspace; added zero-context padded-row behavior. | +| `decode_schedule.py` | `lightllm/common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding.py` | Rewritten around immutable config and caller-owned workspace; removed `infer_state`, global tuning, device inspection, and forward allocation. | +| `copy_latent.py` | `lightllm/common/basemodel/triton_kernel/kv_copy/mla_copy_kv.py` | Added `slot_mapping < 0` padding mask, bounds protection, layout checks, and opt-in-once slot validation. | +| `gather_latent.py` | `lightllm/models/deepseek2/triton_kernel/sample_kv.py` | Replaced modulo sampling with a padding-safe, ragged, request-indirected full-history gather into explicit packed outputs. | + +## Contract + +- Query latent: `[batch, local_heads, 512]`, BF16. +- Query RoPE: `[batch, local_heads, 64]`, BF16. +- Persistent caches: `[slots, 1, 512]` and `[slots, 1, 64]`, BF16. +- Slot/request/context metadata: INT32 CUDA tensors. +- Decode scale: `256 ** -0.5`, matching the pre-absorption GLM QK head + dimension. +- Decode workspaces and outputs are allocated by the caller. The run path does + not inspect device properties or allocate tensors. + +`DEFAULT_GLM_MLA_DECODE_CONFIG` is a conservative correctness configuration +for the initial H100 implementation. It is not recorded as tuned until the +target GLM shape has dedicated benchmark evidence. + +Synchronous slot-value validation is exposed separately so the cache manager +can validate a mapping once and reuse it across layers. Kernel calls remain +bounds-safe; passing `validate_slots=False` or `validate_metadata=False` means +the owning runtime boundary has already established the corresponding value +invariants. diff --git a/src/sparsevllm/triton_kernel/mla/__init__.py b/src/sparsevllm/triton_kernel/mla/__init__.py new file mode 100644 index 00000000..b38d8245 --- /dev/null +++ b/src/sparsevllm/triton_kernel/mla/__init__.py @@ -0,0 +1,40 @@ +"""Latent MLA Triton kernels with explicit Sparse-vLLM contracts.""" + +from .copy_latent import copy_latent_to_cache, validate_copy_slot_mapping +from .decode_schedule import ( + DEFAULT_GLM_MLA_DECODE_CONFIG, + GLM_MLA_SOFTMAX_SCALE, + MlaDecodeLaunchConfig, + MlaDecodeWorkspace, + allocate_mla_decode_workspace, + prepare_mla_decode_schedule, + required_workspace_blocks, + run_mla_decode, + validate_mla_decode_metadata, +) +from .decode_stage1 import MLA_LATENT_DIM, MLA_ROPE_DIM, decode_stage1 +from .decode_stage2 import decode_stage2 +from .gather_latent import ( + gather_latent_history, + validate_gather_metadata, +) + +__all__ = [ + "DEFAULT_GLM_MLA_DECODE_CONFIG", + "GLM_MLA_SOFTMAX_SCALE", + "MLA_LATENT_DIM", + "MLA_ROPE_DIM", + "MlaDecodeLaunchConfig", + "MlaDecodeWorkspace", + "allocate_mla_decode_workspace", + "copy_latent_to_cache", + "decode_stage1", + "decode_stage2", + "gather_latent_history", + "prepare_mla_decode_schedule", + "required_workspace_blocks", + "run_mla_decode", + "validate_copy_slot_mapping", + "validate_gather_metadata", + "validate_mla_decode_metadata", +] diff --git a/src/sparsevllm/triton_kernel/mla/copy_latent.py b/src/sparsevllm/triton_kernel/mla/copy_latent.py new file mode 100644 index 00000000..286ed9c0 --- /dev/null +++ b/src/sparsevllm/triton_kernel/mla/copy_latent.py @@ -0,0 +1,219 @@ +# SPDX-License-Identifier: Apache-2.0 +# Derived from ModelTC/lightllm at commit +# 65c174ee95ac6a6fd36b18b63d0b33d97e76b770: +# lightllm/common/basemodel/triton_kernel/kv_copy/mla_copy_kv.py +# Local changes: mask padded slot_mapping=-1, guard cache bounds, enforce the +# GLM latent layout, and offer explicit duplicate/range validation. + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from .decode_stage1 import MLA_LATENT_DIM, MLA_ROPE_DIM + + +@triton.jit +def _copy_latent_kernel( + latent, + rope, + slot_mapping, + latent_cache, + rope_cache, + stride_latent_token, + stride_latent_head, + stride_latent_dim, + stride_rope_token, + stride_rope_head, + stride_rope_dim, + stride_cache_latent_slot, + stride_cache_latent_head, + stride_cache_latent_dim, + stride_cache_rope_slot, + stride_cache_rope_head, + stride_cache_rope_dim, + cache_slot_count, + LATENT_DIM: tl.constexpr, + ROPE_DIM: tl.constexpr, +): + token_index = tl.program_id(0) + latent_offsets = tl.arange(0, LATENT_DIM) + rope_offsets = tl.arange(0, ROPE_DIM) + destination_slot = tl.load(slot_mapping + token_index).to(tl.int64) + valid_slot = (destination_slot >= 0) & ( + destination_slot < cache_slot_count + ) + safe_slot = tl.where(valid_slot, destination_slot, 0) + + latent_offsets_in = ( + token_index * stride_latent_token + + latent_offsets * stride_latent_dim + ) + rope_offsets_in = ( + token_index * stride_rope_token + rope_offsets * stride_rope_dim + ) + latent_values = tl.load( + latent + latent_offsets_in, + mask=valid_slot, + other=0.0, + ) + rope_values = tl.load( + rope + rope_offsets_in, + mask=valid_slot, + other=0.0, + ) + + latent_offsets_out = ( + safe_slot * stride_cache_latent_slot + + latent_offsets * stride_cache_latent_dim + ) + rope_offsets_out = ( + safe_slot * stride_cache_rope_slot + + rope_offsets * stride_cache_rope_dim + ) + tl.store( + latent_cache + latent_offsets_out, + latent_values, + mask=valid_slot, + ) + tl.store( + rope_cache + rope_offsets_out, + rope_values, + mask=valid_slot, + ) + + +def _validate_copy_tensors( + latent: torch.Tensor, + rope: torch.Tensor, + slot_mapping: torch.Tensor, + latent_cache: torch.Tensor, + rope_cache: torch.Tensor, +) -> None: + tensors = { + "latent": latent, + "rope": rope, + "slot_mapping": slot_mapping, + "latent_cache": latent_cache, + "rope_cache": rope_cache, + } + for name, tensor in tensors.items(): + if tensor.device.type != "cuda": + raise ValueError(f"{name} must be a CUDA tensor, got {tensor.device}") + if tensor.device != latent.device: + raise ValueError( + f"{name} is on {tensor.device}, expected {latent.device}" + ) + for name in ("latent", "rope", "latent_cache", "rope_cache"): + if tensors[name].dtype != torch.bfloat16: + raise TypeError( + f"{name} must use {torch.bfloat16}, got {tensors[name].dtype}" + ) + if slot_mapping.dtype != torch.int32: + raise TypeError( + f"slot_mapping must use {torch.int32}, got {slot_mapping.dtype}" + ) + + if latent.ndim != 3 or latent.shape[1:] != (1, MLA_LATENT_DIM): + raise ValueError( + "latent must have shape [tokens, 1, 512], got " + f"{tuple(latent.shape)}" + ) + if rope.ndim != 3 or rope.shape[1:] != (1, MLA_ROPE_DIM): + raise ValueError( + "rope must have shape [tokens, 1, 64], got " + f"{tuple(rope.shape)}" + ) + if latent_cache.ndim != 3 or latent_cache.shape[1:] != ( + 1, + MLA_LATENT_DIM, + ): + raise ValueError( + "latent_cache must have shape [slots, 1, 512], got " + f"{tuple(latent_cache.shape)}" + ) + if rope_cache.ndim != 3 or rope_cache.shape[1:] != (1, MLA_ROPE_DIM): + raise ValueError( + "rope_cache must have shape [slots, 1, 64], got " + f"{tuple(rope_cache.shape)}" + ) + token_count = latent.shape[0] + if rope.shape[0] != token_count or slot_mapping.shape != (token_count,): + raise ValueError( + "latent, rope, and slot_mapping token dimensions must match" + ) + if latent_cache.shape[0] != rope_cache.shape[0]: + raise ValueError("latent_cache and rope_cache must have equal slots") + + +def validate_copy_slot_mapping( + slot_mapping: torch.Tensor, + *, + cache_slot_count: int, +) -> None: + """Synchronously validate slots once at an owning runtime boundary.""" + + if cache_slot_count <= 0: + raise ValueError("cache_slot_count must be positive") + valid_slots = slot_mapping[slot_mapping >= 0] + if valid_slots.numel() == 0: + return + if bool(torch.any(valid_slots >= cache_slot_count).item()): + raise ValueError( + f"slot_mapping contains a slot outside [0, {cache_slot_count})" + ) + if valid_slots.unique().numel() != valid_slots.numel(): + raise ValueError("slot_mapping contains duplicate non-padding slots") + + +@torch.no_grad() +def copy_latent_to_cache( + latent: torch.Tensor, + rope: torch.Tensor, + slot_mapping: torch.Tensor, + latent_cache: torch.Tensor, + rope_cache: torch.Tensor, + *, + validate_slots: bool = True, +) -> None: + """Copy one token batch into latent MLA caches. + + Negative slots are padding and are never read or written. Set + ``validate_slots=False`` only when the owning cache manager has already + validated this exact mapping; doing so avoids a per-layer host sync. + """ + + _validate_copy_tensors( + latent, + rope, + slot_mapping, + latent_cache, + rope_cache, + ) + token_count = slot_mapping.numel() + if token_count == 0: + return + cache_slot_count = latent_cache.shape[0] + if validate_slots: + validate_copy_slot_mapping( + slot_mapping, + cache_slot_count=cache_slot_count, + ) + + _copy_latent_kernel[(token_count,)]( + latent, + rope, + slot_mapping, + latent_cache, + rope_cache, + *latent.stride(), + *rope.stride(), + *latent_cache.stride(), + *rope_cache.stride(), + cache_slot_count=cache_slot_count, + LATENT_DIM=MLA_LATENT_DIM, + ROPE_DIM=MLA_ROPE_DIM, + num_warps=1, + num_stages=1, + ) diff --git a/src/sparsevllm/triton_kernel/mla/decode_schedule.py b/src/sparsevllm/triton_kernel/mla/decode_schedule.py new file mode 100644 index 00000000..09d563f9 --- /dev/null +++ b/src/sparsevllm/triton_kernel/mla/decode_schedule.py @@ -0,0 +1,398 @@ +# SPDX-License-Identifier: Apache-2.0 +# Scheduling math derived from ModelTC/lightllm at commit +# 65c174ee95ac6a6fd36b18b63d0b33d97e76b770: +# lightllm/common/basemodel/triton_kernel/mla_att/decode_att/ +# gqa_flash_decoding.py +# Local rewrite: caller-owned workspace, immutable launch configuration, no +# infer_state/global config/device probing, and no allocation in the run path. + +from __future__ import annotations + +from dataclasses import dataclass + +import torch +import triton +import triton.language as tl + +from .decode_stage1 import MLA_LATENT_DIM, decode_stage1 +from .decode_stage2 import decode_stage2 + + +GLM_MLA_SOFTMAX_SCALE = 256**-0.5 + + +@dataclass(frozen=True, slots=True) +class MlaDecodeLaunchConfig: + """Static launch configuration for the two-stage decode kernel.""" + + program_count: int = 128 + blocks_per_program: int = 4 + block_n: int = 16 + block_q_heads: int = 16 + stage1_num_warps: int = 4 + stage1_pipeline_stages: int = 2 + stage2_num_warps: int = 4 + stage2_pipeline_stages: int = 2 + + def __post_init__(self) -> None: + positive_fields = ( + "program_count", + "blocks_per_program", + "block_n", + "block_q_heads", + "stage1_pipeline_stages", + "stage2_pipeline_stages", + ) + for field_name in positive_fields: + if getattr(self, field_name) <= 0: + raise ValueError(f"{field_name} must be positive") + for field_name in ("block_n", "block_q_heads"): + value = getattr(self, field_name) + if value & (value - 1): + raise ValueError(f"{field_name} must be a power of two") + for field_name in ("stage1_num_warps", "stage2_num_warps"): + if getattr(self, field_name) not in {1, 2, 4, 8}: + raise ValueError(f"{field_name} must be one of 1, 2, 4, or 8") + + +DEFAULT_GLM_MLA_DECODE_CONFIG = MlaDecodeLaunchConfig() + + +@dataclass(frozen=True, slots=True) +class MlaDecodeWorkspace: + """Caller-owned tensors required by MLA decode.""" + + block_size: torch.Tensor + batch_start_indices: torch.Tensor + mid_output: torch.Tensor + mid_logsumexp: torch.Tensor + + +def required_workspace_blocks( + batch_size: int, + config: MlaDecodeLaunchConfig, +) -> int: + if batch_size <= 0: + raise ValueError("batch_size must be positive") + return config.program_count * config.blocks_per_program + batch_size + + +def allocate_mla_decode_workspace( + *, + batch_size: int, + head_count: int, + device: torch.device | str, + config: MlaDecodeLaunchConfig = DEFAULT_GLM_MLA_DECODE_CONFIG, +) -> MlaDecodeWorkspace: + """Allocate decode workspace outside the attention execution path.""" + + if head_count <= 0: + raise ValueError("head_count must be positive") + block_capacity = required_workspace_blocks(batch_size, config) + return MlaDecodeWorkspace( + block_size=torch.empty((1,), dtype=torch.int32, device=device), + batch_start_indices=torch.empty( + (batch_size,), + dtype=torch.int32, + device=device, + ), + mid_output=torch.empty( + (head_count, block_capacity, MLA_LATENT_DIM), + dtype=torch.float32, + device=device, + ), + mid_logsumexp=torch.empty( + (head_count, block_capacity), + dtype=torch.float32, + device=device, + ), + ) + + +@triton.jit +def _build_decode_schedule_kernel( + context_lens, + block_size_ptr, + batch_start_indices, + program_count, + blocks_per_program, + batch_size, + BLOCK_N: tl.constexpr, + PADDED_BATCH_SIZE: tl.constexpr, +): + offsets = tl.arange(0, PADDED_BATCH_SIZE) + context_mask = offsets < batch_size + lengths = tl.load( + context_lens + offsets, + mask=context_mask, + other=0, + ) + total_tokens = tl.sum(lengths, axis=0) + target_blocks = program_count * blocks_per_program + unaligned_block_size = tl.maximum(1, tl.cdiv(total_tokens, target_blocks)) + block_size = tl.cdiv(unaligned_block_size, BLOCK_N) * BLOCK_N + + block_counts = tl.cdiv(lengths, block_size) + cumulative_blocks = tl.cumsum(block_counts, axis=0) + starts = cumulative_blocks - block_counts + tl.store( + batch_start_indices + offsets, + starts, + mask=context_mask, + ) + tl.store(block_size_ptr, block_size) + + +def _validate_schedule_workspace( + context_lens: torch.Tensor, + workspace: MlaDecodeWorkspace, + config: MlaDecodeLaunchConfig, +) -> None: + if context_lens.device.type != "cuda": + raise ValueError( + f"context_lens must be a CUDA tensor, got {context_lens.device}" + ) + if context_lens.dtype != torch.int32: + raise TypeError( + f"context_lens must use {torch.int32}, got {context_lens.dtype}" + ) + if context_lens.ndim != 1 or context_lens.numel() == 0: + raise ValueError("context_lens must be a non-empty one-dimensional tensor") + + batch_size = context_lens.numel() + workspace_tensors = { + "block_size": workspace.block_size, + "batch_start_indices": workspace.batch_start_indices, + "mid_output": workspace.mid_output, + "mid_logsumexp": workspace.mid_logsumexp, + } + for name, tensor in workspace_tensors.items(): + if tensor.device != context_lens.device: + raise ValueError( + f"workspace {name} is on {tensor.device}, expected " + f"{context_lens.device}" + ) + if workspace.block_size.dtype != torch.int32 or workspace.block_size.shape != ( + 1, + ): + raise ValueError("workspace.block_size must be int32 with shape [1]") + if ( + workspace.batch_start_indices.dtype != torch.int32 + or workspace.batch_start_indices.ndim != 1 + or workspace.batch_start_indices.numel() < batch_size + ): + raise ValueError( + "workspace.batch_start_indices must be int32 and have capacity " + f"for {batch_size} rows" + ) + required_blocks = required_workspace_blocks(batch_size, config) + if ( + workspace.mid_output.dtype != torch.float32 + or workspace.mid_output.ndim != 3 + or workspace.mid_output.shape[-1] != MLA_LATENT_DIM + ): + raise ValueError( + "workspace.mid_output must be float32 with shape " + "[heads, blocks, 512]" + ) + if workspace.mid_output.shape[1] < required_blocks: + raise ValueError( + "MLA decode workspace is too small: " + f"blocks={workspace.mid_output.shape[1]}, " + f"required={required_blocks}" + ) + if ( + workspace.mid_logsumexp.dtype != torch.float32 + or workspace.mid_logsumexp.shape != workspace.mid_output.shape[:2] + ): + raise ValueError( + "workspace.mid_logsumexp must be float32 and match the first " + "two mid_output dimensions" + ) + + +@torch.no_grad() +def prepare_mla_decode_schedule( + context_lens: torch.Tensor, + workspace: MlaDecodeWorkspace, + *, + config: MlaDecodeLaunchConfig = DEFAULT_GLM_MLA_DECODE_CONFIG, +) -> None: + """Fill device-side block size and batch offsets without a CPU sync.""" + + _validate_schedule_workspace(context_lens, workspace, config) + batch_size = context_lens.numel() + _build_decode_schedule_kernel[(1,)]( + context_lens, + workspace.block_size, + workspace.batch_start_indices, + program_count=config.program_count, + blocks_per_program=config.blocks_per_program, + batch_size=batch_size, + BLOCK_N=config.block_n, + PADDED_BATCH_SIZE=triton.next_power_of_2(batch_size), + num_warps=4, + num_stages=1, + ) + + +def validate_mla_decode_metadata( + active_slots: torch.Tensor, + request_indices: torch.Tensor, + context_lens: torch.Tensor, + *, + cache_slot_count: int, +) -> None: + """Synchronously validate one decode view before per-layer reuse.""" + + metadata = { + "active_slots": active_slots, + "request_indices": request_indices, + "context_lens": context_lens, + } + for name, tensor in metadata.items(): + if tensor.device.type != "cuda": + raise ValueError(f"{name} must be a CUDA tensor, got {tensor.device}") + if tensor.device != context_lens.device: + raise ValueError( + f"{name} is on {tensor.device}, expected {context_lens.device}" + ) + if tensor.dtype != torch.int32: + raise TypeError( + f"{name} must use {torch.int32}, got {tensor.dtype}" + ) + if active_slots.ndim != 2: + raise ValueError("active_slots must have shape [rows, max_context_len]") + if context_lens.ndim != 1 or context_lens.numel() == 0: + raise ValueError("context_lens must be a non-empty one-dimensional tensor") + batch_size = context_lens.numel() + if request_indices.shape != (batch_size,): + raise ValueError( + f"request_indices must have shape ({batch_size},), got " + f"{tuple(request_indices.shape)}" + ) + if cache_slot_count <= 0: + raise ValueError("cache_slot_count must be positive") + + request_rows = request_indices.tolist() + lengths = context_lens.tolist() + real_request_rows: list[int] = [] + for batch_index, (request_row, length) in enumerate( + zip(request_rows, lengths) + ): + if length < 0 or length > active_slots.shape[1]: + raise ValueError( + f"context_lens[{batch_index}]={length} is outside " + f"[0, {active_slots.shape[1]}]" + ) + if request_row < 0: + if length != 0: + raise ValueError( + "padded request rows must have zero context length" + ) + continue + if request_row >= active_slots.shape[0]: + raise ValueError( + f"request_indices[{batch_index}]={request_row} is outside " + f"[0, {active_slots.shape[0]})" + ) + real_request_rows.append(request_row) + if length == 0: + continue + slots = active_slots[request_row, :length] + if bool(torch.any(slots < 0).item()) or bool( + torch.any(slots >= cache_slot_count).item() + ): + raise ValueError( + f"active_slots row {request_row} contains an invalid slot" + ) + if slots.unique().numel() != slots.numel(): + raise ValueError( + f"active_slots row {request_row} contains duplicate slots" + ) + if len(set(real_request_rows)) != len(real_request_rows): + raise ValueError("request_indices contains duplicate non-padding rows") + + +@torch.no_grad() +def run_mla_decode( + q_latent: torch.Tensor, + q_rope: torch.Tensor, + latent_cache: torch.Tensor, + rope_cache: torch.Tensor, + active_slots: torch.Tensor, + request_indices: torch.Tensor, + context_lens: torch.Tensor, + output: torch.Tensor, + workspace: MlaDecodeWorkspace, + *, + config: MlaDecodeLaunchConfig = DEFAULT_GLM_MLA_DECODE_CONFIG, + validate_metadata: bool = True, +) -> torch.Tensor: + """Run GLM MLA decode using only explicit tensors and static config.""" + + _validate_schedule_workspace(context_lens, workspace, config) + if q_latent.ndim != 3: + raise ValueError("q_latent must have shape [batch, heads, 512]") + batch_size, head_count = q_latent.shape[:2] + if context_lens.numel() != batch_size: + raise ValueError( + "q_latent batch size and context_lens length must match: " + f"{batch_size} != {context_lens.numel()}" + ) + if output.shape != q_latent.shape: + raise ValueError( + f"output must have shape {tuple(q_latent.shape)}, got " + f"{tuple(output.shape)}" + ) + if output.device != q_latent.device: + raise ValueError( + f"output is on {output.device}, expected {q_latent.device}" + ) + if output.dtype != torch.bfloat16: + raise TypeError( + f"output must use {torch.bfloat16}, got {output.dtype}" + ) + if workspace.mid_output.shape[0] < head_count: + raise ValueError( + "MLA decode workspace is too small: " + f"heads={workspace.mid_output.shape[0]}, required={head_count}" + ) + if validate_metadata: + validate_mla_decode_metadata( + active_slots, + request_indices, + context_lens, + cache_slot_count=latent_cache.shape[0], + ) + + prepare_mla_decode_schedule(context_lens, workspace, config=config) + decode_stage1( + q_latent, + q_rope, + latent_cache, + rope_cache, + active_slots, + request_indices, + context_lens, + workspace.block_size, + workspace.mid_output, + workspace.mid_logsumexp, + softmax_scale=GLM_MLA_SOFTMAX_SCALE, + program_count=config.program_count, + block_q_heads=config.block_q_heads, + block_n=config.block_n, + pipeline_stages=config.stage1_pipeline_stages, + num_warps=config.stage1_num_warps, + ) + decode_stage2( + workspace.block_size, + workspace.batch_start_indices, + context_lens, + workspace.mid_output, + workspace.mid_logsumexp, + output, + pipeline_stages=config.stage2_pipeline_stages, + num_warps=config.stage2_num_warps, + ) + return output diff --git a/src/sparsevllm/triton_kernel/mla/decode_stage1.py b/src/sparsevllm/triton_kernel/mla/decode_stage1.py new file mode 100644 index 00000000..55702ffe --- /dev/null +++ b/src/sparsevllm/triton_kernel/mla/decode_stage1.py @@ -0,0 +1,391 @@ +# SPDX-License-Identifier: Apache-2.0 +# Derived from ModelTC/lightllm at commit +# 65c174ee95ac6a6fd36b18b63d0b33d97e76b770: +# lightllm/common/basemodel/triton_kernel/mla_att/decode_att/ +# gqa_flash_decoding_stage1.py +# Local changes: remove LightLLM runtime/device helpers, expose an explicit +# workspace API, restrict the layout to the GLM MLA contract, and preserve +# arbitrary tensor strides. + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +MLA_LATENT_DIM = 512 +MLA_ROPE_DIM = 64 + + +@triton.jit +def _decode_stage1_kernel( + q_latent, + q_rope, + latent_cache, + rope_cache, + softmax_scale, + active_slots, + request_indices, + context_lens, + mid_output, + mid_logsumexp, + stride_slots_row, + stride_slots_token, + stride_q_latent_batch, + stride_q_latent_head, + stride_q_latent_dim, + stride_q_rope_batch, + stride_q_rope_head, + stride_q_rope_dim, + stride_latent_slot, + stride_latent_head, + stride_latent_dim, + stride_rope_slot, + stride_rope_head, + stride_rope_dim, + stride_mid_head, + stride_mid_block, + stride_mid_dim, + stride_lse_head, + stride_lse_block, + block_size_ptr, + cache_slot_count, + program_count, + head_group_count, + head_count, + batch_size, + LATENT_DIM: tl.constexpr, + ROPE_DIM: tl.constexpr, + BLOCK_Q_HEADS: tl.constexpr, + BLOCK_N: tl.constexpr, + PIPELINE_STAGES: tl.constexpr, + MASK_HEADS: tl.constexpr, +): + program_id = tl.program_id(0).to(tl.int64) + output_batch_start = tl.cast(0, tl.int64) + block_size = tl.load(block_size_ptr, eviction_policy="evict_last") + + head_offsets = tl.arange(0, BLOCK_Q_HEADS) + latent_offsets = tl.arange(0, LATENT_DIM) + rope_offsets = tl.arange(0, ROPE_DIM) + + for batch_index in range(batch_size): + context_len = tl.load( + context_lens + batch_index, + eviction_policy="evict_last", + ) + block_count = tl.cdiv(context_len, block_size) + work_count = block_count * head_group_count + request_index = tl.load( + request_indices + batch_index, + eviction_policy="evict_last", + ) + slots_row = active_slots + request_index * stride_slots_row + + work_index = program_id + while work_index < work_count: + head_group_index = work_index % head_group_count + sequence_block_index = work_index // head_group_count + query_heads = head_group_index * BLOCK_Q_HEADS + head_offsets + if MASK_HEADS: + head_mask = query_heads < head_count + + block_start = block_size * sequence_block_index + block_end = tl.minimum(context_len, block_start + block_size) + + q_latent_offsets = ( + batch_index * stride_q_latent_batch + + query_heads[:, None] * stride_q_latent_head + + latent_offsets[None, :] * stride_q_latent_dim + ) + q_rope_offsets = ( + batch_index * stride_q_rope_batch + + query_heads[:, None] * stride_q_rope_head + + rope_offsets[None, :] * stride_q_rope_dim + ) + if MASK_HEADS: + query_latent = tl.load( + q_latent + q_latent_offsets, + mask=head_mask[:, None], + other=0.0, + ) + query_rope = tl.load( + q_rope + q_rope_offsets, + mask=head_mask[:, None], + other=0.0, + ) + else: + query_latent = tl.load(q_latent + q_latent_offsets) + query_rope = tl.load(q_rope + q_rope_offsets) + + loop_count = tl.cdiv(block_end - block_start, BLOCK_N) + token_offsets = block_start + tl.arange(0, BLOCK_N) + sum_exp = tl.zeros([BLOCK_Q_HEADS], dtype=tl.float32) + max_logit = tl.full( + [BLOCK_Q_HEADS], + -float("inf"), + dtype=tl.float32, + ) + accumulator = tl.zeros( + [BLOCK_Q_HEADS, LATENT_DIM], + dtype=tl.float32, + ) + + for token_block in tl.range( + 0, + loop_count, + 1, + num_stages=PIPELINE_STAGES, + ): + token_indices = token_block * BLOCK_N + token_offsets + token_mask = token_indices < block_end + cache_slots = tl.load( + slots_row + token_indices * stride_slots_token, + mask=token_mask, + other=0, + ).to(tl.int64) + valid_cache_slot = ( + token_mask + & (cache_slots >= 0) + & (cache_slots < cache_slot_count) + ) + safe_cache_slots = tl.where(valid_cache_slot, cache_slots, 0) + + latent_cache_offsets = ( + safe_cache_slots[None, :] * stride_latent_slot + + latent_offsets[:, None] * stride_latent_dim + ) + cached_latent = tl.load( + latent_cache + latent_cache_offsets, + mask=valid_cache_slot[None, :], + other=0.0, + ) + logits = tl.dot(query_latent, cached_latent) + + rope_cache_offsets = ( + safe_cache_slots[None, :] * stride_rope_slot + + rope_offsets[:, None] * stride_rope_dim + ) + cached_rope = tl.load( + rope_cache + rope_cache_offsets, + mask=valid_cache_slot[None, :], + other=0.0, + ) + logits += tl.dot(query_rope, cached_rope) + logits *= softmax_scale + logits = tl.where( + valid_cache_slot[None, :], + logits, + -float("inf"), + ) + + block_max = tl.max(logits, axis=1) + new_max = tl.maximum(block_max, max_logit) + exp_logits = tl.exp(logits - new_max[:, None]) + old_scale = tl.exp(max_logit - new_max) + accumulator *= old_scale[:, None] + accumulator += tl.dot( + exp_logits.to(cached_latent.dtype), + tl.trans(cached_latent), + ) + sum_exp = sum_exp * old_scale + tl.sum(exp_logits, axis=1) + max_logit = new_max + + output_block_index = output_batch_start + sequence_block_index + mid_offsets = ( + query_heads[:, None] * stride_mid_head + + output_block_index * stride_mid_block + + latent_offsets[None, :] * stride_mid_dim + ) + lse_offsets = ( + query_heads * stride_lse_head + + output_block_index * stride_lse_block + ) + normalized = accumulator / sum_exp[:, None] + logsumexp = max_logit + tl.log(sum_exp) + if MASK_HEADS: + tl.store( + mid_output + mid_offsets, + normalized, + mask=head_mask[:, None], + ) + tl.store( + mid_logsumexp + lse_offsets, + logsumexp, + mask=head_mask, + ) + else: + tl.store(mid_output + mid_offsets, normalized) + tl.store(mid_logsumexp + lse_offsets, logsumexp) + + work_index += program_count + + output_batch_start += block_count + + +def _require_cuda_tensor(name: str, tensor: torch.Tensor) -> None: + if tensor.device.type != "cuda": + raise ValueError(f"{name} must be a CUDA tensor, got {tensor.device}") + + +def _require_dtype( + name: str, + tensor: torch.Tensor, + expected: torch.dtype, +) -> None: + if tensor.dtype != expected: + raise TypeError(f"{name} must use {expected}, got {tensor.dtype}") + + +@torch.no_grad() +def decode_stage1( + q_latent: torch.Tensor, + q_rope: torch.Tensor, + latent_cache: torch.Tensor, + rope_cache: torch.Tensor, + active_slots: torch.Tensor, + request_indices: torch.Tensor, + context_lens: torch.Tensor, + block_size: torch.Tensor, + mid_output: torch.Tensor, + mid_logsumexp: torch.Tensor, + *, + softmax_scale: float, + program_count: int, + block_q_heads: int, + block_n: int, + pipeline_stages: int, + num_warps: int, +) -> None: + """Compute independently normalized MLA attention blocks. + + All scheduling tensors and workspaces are caller-owned. The function does + not allocate, inspect device properties, or depend on model/runtime state. + """ + + tensors = { + "q_latent": q_latent, + "q_rope": q_rope, + "latent_cache": latent_cache, + "rope_cache": rope_cache, + "active_slots": active_slots, + "request_indices": request_indices, + "context_lens": context_lens, + "block_size": block_size, + "mid_output": mid_output, + "mid_logsumexp": mid_logsumexp, + } + for name, tensor in tensors.items(): + _require_cuda_tensor(name, tensor) + if tensor.device != q_latent.device: + raise ValueError( + f"{name} is on {tensor.device}, expected {q_latent.device}" + ) + + for name in ("q_latent", "q_rope", "latent_cache", "rope_cache"): + _require_dtype(name, tensors[name], torch.bfloat16) + for name in ("active_slots", "request_indices", "context_lens"): + _require_dtype(name, tensors[name], torch.int32) + _require_dtype("block_size", block_size, torch.int32) + _require_dtype("mid_output", mid_output, torch.float32) + _require_dtype("mid_logsumexp", mid_logsumexp, torch.float32) + + if q_latent.ndim != 3 or q_latent.shape[-1] != MLA_LATENT_DIM: + raise ValueError( + "q_latent must have shape [batch, heads, 512], got " + f"{tuple(q_latent.shape)}" + ) + if q_rope.shape != (*q_latent.shape[:-1], MLA_ROPE_DIM): + raise ValueError( + "q_rope must have shape [batch, heads, 64], got " + f"{tuple(q_rope.shape)}" + ) + if latent_cache.ndim != 3 or latent_cache.shape[1:] != ( + 1, + MLA_LATENT_DIM, + ): + raise ValueError( + "latent_cache must have shape [slots, 1, 512], got " + f"{tuple(latent_cache.shape)}" + ) + if rope_cache.ndim != 3 or rope_cache.shape[1:] != (1, MLA_ROPE_DIM): + raise ValueError( + "rope_cache must have shape [slots, 1, 64], got " + f"{tuple(rope_cache.shape)}" + ) + if latent_cache.shape[0] != rope_cache.shape[0]: + raise ValueError("latent_cache and rope_cache must have equal slots") + if active_slots.ndim != 2: + raise ValueError("active_slots must have shape [rows, max_context_len]") + + batch_size, head_count = q_latent.shape[:2] + if request_indices.shape != (batch_size,): + raise ValueError( + f"request_indices must have shape ({batch_size},), got " + f"{tuple(request_indices.shape)}" + ) + if context_lens.shape != (batch_size,): + raise ValueError( + f"context_lens must have shape ({batch_size},), got " + f"{tuple(context_lens.shape)}" + ) + if block_size.shape != (1,): + raise ValueError("block_size must have shape [1]") + if mid_output.ndim != 3 or mid_output.shape[2] != MLA_LATENT_DIM: + raise ValueError("mid_output must have shape [heads, blocks, 512]") + if mid_logsumexp.shape != mid_output.shape[:2]: + raise ValueError( + "mid_logsumexp must match the first two mid_output dimensions" + ) + if mid_output.shape[0] < head_count: + raise ValueError( + f"mid_output has capacity for {mid_output.shape[0]} heads, " + f"but {head_count} are required" + ) + if program_count <= 0: + raise ValueError("program_count must be positive") + if block_q_heads <= 0 or block_q_heads & (block_q_heads - 1): + raise ValueError("block_q_heads must be a positive power of two") + if block_n <= 0 or block_n & (block_n - 1): + raise ValueError("block_n must be a positive power of two") + if pipeline_stages <= 0: + raise ValueError("pipeline_stages must be positive") + if num_warps not in {1, 2, 4, 8}: + raise ValueError("num_warps must be one of 1, 2, 4, or 8") + + head_group_count = triton.cdiv(head_count, block_q_heads) + mask_heads = head_count % block_q_heads != 0 + _decode_stage1_kernel[(program_count,)]( + q_latent, + q_rope, + latent_cache, + rope_cache, + softmax_scale, + active_slots, + request_indices, + context_lens, + mid_output, + mid_logsumexp, + *active_slots.stride(), + *q_latent.stride(), + *q_rope.stride(), + *latent_cache.stride(), + *rope_cache.stride(), + *mid_output.stride(), + *mid_logsumexp.stride(), + block_size, + cache_slot_count=latent_cache.shape[0], + program_count=program_count, + head_group_count=head_group_count, + head_count=head_count, + batch_size=batch_size, + LATENT_DIM=MLA_LATENT_DIM, + ROPE_DIM=MLA_ROPE_DIM, + BLOCK_Q_HEADS=block_q_heads, + BLOCK_N=block_n, + PIPELINE_STAGES=pipeline_stages, + MASK_HEADS=mask_heads, + num_warps=num_warps, + num_stages=1, + ) diff --git a/src/sparsevllm/triton_kernel/mla/decode_stage2.py b/src/sparsevllm/triton_kernel/mla/decode_stage2.py new file mode 100644 index 00000000..b07aa5cf --- /dev/null +++ b/src/sparsevllm/triton_kernel/mla/decode_stage2.py @@ -0,0 +1,175 @@ +# SPDX-License-Identifier: Apache-2.0 +# Derived from ModelTC/lightllm at commit +# 65c174ee95ac6a6fd36b18b63d0b33d97e76b770: +# lightllm/common/basemodel/triton_kernel/mla_att/decode_att/ +# gqa_flash_decoding_stage2.py +# Local changes: remove unused runtime imports, expose explicit scheduling and +# workspace tensors, preserve arbitrary strides, and restrict dimensions to +# the GLM MLA contract. + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from .decode_stage1 import MLA_LATENT_DIM + + +@triton.jit +def _decode_stage2_kernel( + block_size_ptr, + batch_start_indices, + context_lens, + mid_output, + mid_logsumexp, + output, + stride_mid_head, + stride_mid_block, + stride_mid_dim, + stride_lse_head, + stride_lse_block, + stride_output_batch, + stride_output_head, + stride_output_dim, + LATENT_DIM: tl.constexpr, + PIPELINE_STAGES: tl.constexpr, +): + head_index = tl.program_id(0) + batch_index = tl.program_id(1) + dim_offsets = tl.arange(0, LATENT_DIM) + + context_len = tl.load(context_lens + batch_index) + batch_start = tl.load(batch_start_indices + batch_index) + block_size = tl.load(block_size_ptr) + block_count = tl.cdiv(context_len, block_size) + + sum_exp = 0.0 + max_logit = -float("inf") + accumulator = tl.zeros([LATENT_DIM], dtype=tl.float32) + mid_offsets = ( + head_index * stride_mid_head + + batch_start * stride_mid_block + + dim_offsets * stride_mid_dim + ) + lse_offset = ( + head_index * stride_lse_head + + batch_start * stride_lse_block + ) + + for block_index in tl.range( + 0, + block_count, + 1, + num_stages=PIPELINE_STAGES, + ): + block_output = tl.load( + mid_output + mid_offsets + block_index * stride_mid_block + ) + block_lse = tl.load( + mid_logsumexp + lse_offset + block_index * stride_lse_block + ) + new_max = tl.maximum(block_lse, max_logit) + old_scale = tl.exp(max_logit - new_max) + block_scale = tl.exp(block_lse - new_max) + accumulator = accumulator * old_scale + block_scale * block_output + sum_exp = sum_exp * old_scale + block_scale + max_logit = new_max + + output_offsets = ( + batch_index * stride_output_batch + + head_index * stride_output_head + + dim_offsets * stride_output_dim + ) + normalized = tl.where(block_count > 0, accumulator / sum_exp, 0.0) + tl.store(output + output_offsets, normalized) + + +@torch.no_grad() +def decode_stage2( + block_size: torch.Tensor, + batch_start_indices: torch.Tensor, + context_lens: torch.Tensor, + mid_output: torch.Tensor, + mid_logsumexp: torch.Tensor, + output: torch.Tensor, + *, + pipeline_stages: int, + num_warps: int, +) -> None: + """Merge independently normalized stage-one blocks into final output.""" + + tensors = { + "block_size": block_size, + "batch_start_indices": batch_start_indices, + "context_lens": context_lens, + "mid_output": mid_output, + "mid_logsumexp": mid_logsumexp, + "output": output, + } + for name, tensor in tensors.items(): + if tensor.device.type != "cuda": + raise ValueError(f"{name} must be a CUDA tensor, got {tensor.device}") + if tensor.device != output.device: + raise ValueError( + f"{name} is on {tensor.device}, expected {output.device}" + ) + + for name in ("block_size", "batch_start_indices", "context_lens"): + if tensors[name].dtype != torch.int32: + raise TypeError( + f"{name} must use {torch.int32}, got {tensors[name].dtype}" + ) + if mid_output.dtype != torch.float32: + raise TypeError("mid_output must use torch.float32") + if mid_logsumexp.dtype != torch.float32: + raise TypeError("mid_logsumexp must use torch.float32") + if output.dtype != torch.bfloat16: + raise TypeError("output must use torch.bfloat16") + + if block_size.shape != (1,): + raise ValueError("block_size must have shape [1]") + if output.ndim != 3 or output.shape[-1] != MLA_LATENT_DIM: + raise ValueError("output must have shape [batch, heads, 512]") + batch_size, head_count = output.shape[:2] + if context_lens.shape != (batch_size,): + raise ValueError( + f"context_lens must have shape ({batch_size},), got " + f"{tuple(context_lens.shape)}" + ) + if batch_start_indices.ndim != 1 or batch_start_indices.numel() < batch_size: + raise ValueError( + "batch_start_indices must be one-dimensional with at least " + f"{batch_size} entries" + ) + if mid_output.ndim != 3 or mid_output.shape[-1] != MLA_LATENT_DIM: + raise ValueError("mid_output must have shape [heads, blocks, 512]") + if mid_output.shape[0] < head_count: + raise ValueError( + f"mid_output has capacity for {mid_output.shape[0]} heads, " + f"but {head_count} are required" + ) + if mid_logsumexp.shape != mid_output.shape[:2]: + raise ValueError( + "mid_logsumexp must match the first two mid_output dimensions" + ) + if pipeline_stages <= 0: + raise ValueError("pipeline_stages must be positive") + if num_warps not in {1, 2, 4, 8}: + raise ValueError("num_warps must be one of 1, 2, 4, or 8") + + _decode_stage2_kernel[(head_count, batch_size)]( + block_size, + batch_start_indices, + context_lens, + mid_output, + mid_logsumexp, + output, + *mid_output.stride(), + *mid_logsumexp.stride(), + *output.stride(), + LATENT_DIM=MLA_LATENT_DIM, + PIPELINE_STAGES=pipeline_stages, + num_warps=num_warps, + num_stages=1, + ) diff --git a/src/sparsevllm/triton_kernel/mla/gather_latent.py b/src/sparsevllm/triton_kernel/mla/gather_latent.py new file mode 100644 index 00000000..88bbcfa5 --- /dev/null +++ b/src/sparsevllm/triton_kernel/mla/gather_latent.py @@ -0,0 +1,336 @@ +# SPDX-License-Identifier: Apache-2.0 +# Derived from ModelTC/lightllm at commit +# 65c174ee95ac6a6fd36b18b63d0b33d97e76b770: +# lightllm/models/deepseek2/triton_kernel/sample_kv.py +# Local changes: gather full ragged history without modulo sampling, support +# request indirection/non-contiguous strides/padded rows, and validate packed +# destinations explicitly. + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from .decode_stage1 import MLA_LATENT_DIM, MLA_ROPE_DIM + + +@triton.jit +def _gather_latent_kernel( + latent_cache, + rope_cache, + active_slots, + request_indices, + context_lens, + packed_start_locs, + gathered_latent, + gathered_rope, + stride_latent_slot, + stride_latent_head, + stride_latent_dim, + stride_rope_slot, + stride_rope_head, + stride_rope_dim, + stride_slots_row, + stride_slots_token, + stride_request, + stride_context, + stride_packed_start, + stride_output_latent_token, + stride_output_latent_dim, + stride_output_rope_token, + stride_output_rope_dim, + cache_slot_count, + output_capacity, + LATENT_DIM: tl.constexpr, + ROPE_DIM: tl.constexpr, + BLOCK_SEQ: tl.constexpr, +): + batch_index = tl.program_id(0) + sequence_block = tl.program_id(1) + context_len = tl.load(context_lens + batch_index * stride_context) + request_index = tl.load( + request_indices + batch_index * stride_request + ).to(tl.int64) + packed_start = tl.load( + packed_start_locs + batch_index * stride_packed_start + ).to(tl.int64) + + token_offsets = sequence_block * BLOCK_SEQ + tl.arange(0, BLOCK_SEQ) + valid_row = (request_index >= 0) & (context_len > 0) + valid_token = valid_row & (token_offsets < context_len) + safe_request_index = tl.where(valid_row, request_index, 0) + cache_slots = tl.load( + active_slots + + safe_request_index * stride_slots_row + + token_offsets * stride_slots_token, + mask=valid_token, + other=0, + ).to(tl.int64) + valid_slot = valid_token & (cache_slots >= 0) & ( + cache_slots < cache_slot_count + ) + safe_cache_slots = tl.where(valid_slot, cache_slots, 0) + + latent_offsets = tl.arange(0, LATENT_DIM) + rope_offsets = tl.arange(0, ROPE_DIM) + latent_cache_offsets = ( + safe_cache_slots[:, None] * stride_latent_slot + + latent_offsets[None, :] * stride_latent_dim + ) + rope_cache_offsets = ( + safe_cache_slots[:, None] * stride_rope_slot + + rope_offsets[None, :] * stride_rope_dim + ) + latent_values = tl.load( + latent_cache + latent_cache_offsets, + mask=valid_slot[:, None], + other=0.0, + ) + rope_values = tl.load( + rope_cache + rope_cache_offsets, + mask=valid_slot[:, None], + other=0.0, + ) + + output_tokens = packed_start + token_offsets + valid_output = ( + valid_token + & (output_tokens >= 0) + & (output_tokens < output_capacity) + ) + output_latent_offsets = ( + output_tokens[:, None] * stride_output_latent_token + + latent_offsets[None, :] * stride_output_latent_dim + ) + output_rope_offsets = ( + output_tokens[:, None] * stride_output_rope_token + + rope_offsets[None, :] * stride_output_rope_dim + ) + tl.store( + gathered_latent + output_latent_offsets, + latent_values, + mask=valid_output[:, None], + ) + tl.store( + gathered_rope + output_rope_offsets, + rope_values, + mask=valid_output[:, None], + ) + + +def _validate_gather_tensors( + latent_cache: torch.Tensor, + rope_cache: torch.Tensor, + active_slots: torch.Tensor, + request_indices: torch.Tensor, + context_lens: torch.Tensor, + packed_start_locs: torch.Tensor, + gathered_latent: torch.Tensor, + gathered_rope: torch.Tensor, +) -> None: + tensors = { + "latent_cache": latent_cache, + "rope_cache": rope_cache, + "active_slots": active_slots, + "request_indices": request_indices, + "context_lens": context_lens, + "packed_start_locs": packed_start_locs, + "gathered_latent": gathered_latent, + "gathered_rope": gathered_rope, + } + for name, tensor in tensors.items(): + if tensor.device.type != "cuda": + raise ValueError(f"{name} must be a CUDA tensor, got {tensor.device}") + if tensor.device != latent_cache.device: + raise ValueError( + f"{name} is on {tensor.device}, expected {latent_cache.device}" + ) + for name in ( + "latent_cache", + "rope_cache", + "gathered_latent", + "gathered_rope", + ): + if tensors[name].dtype != torch.bfloat16: + raise TypeError( + f"{name} must use {torch.bfloat16}, got {tensors[name].dtype}" + ) + for name in ( + "active_slots", + "request_indices", + "context_lens", + "packed_start_locs", + ): + if tensors[name].dtype != torch.int32: + raise TypeError( + f"{name} must use {torch.int32}, got {tensors[name].dtype}" + ) + + if latent_cache.ndim != 3 or latent_cache.shape[1:] != ( + 1, + MLA_LATENT_DIM, + ): + raise ValueError("latent_cache must have shape [slots, 1, 512]") + if rope_cache.ndim != 3 or rope_cache.shape[1:] != (1, MLA_ROPE_DIM): + raise ValueError("rope_cache must have shape [slots, 1, 64]") + if latent_cache.shape[0] != rope_cache.shape[0]: + raise ValueError("latent_cache and rope_cache must have equal slots") + if active_slots.ndim != 2: + raise ValueError("active_slots must have shape [rows, max_context_len]") + if context_lens.ndim != 1 or context_lens.numel() == 0: + raise ValueError("context_lens must be a non-empty one-dimensional tensor") + batch_size = context_lens.numel() + if request_indices.shape != (batch_size,): + raise ValueError( + f"request_indices must have shape ({batch_size},), got " + f"{tuple(request_indices.shape)}" + ) + if packed_start_locs.shape != (batch_size,): + raise ValueError( + f"packed_start_locs must have shape ({batch_size},), got " + f"{tuple(packed_start_locs.shape)}" + ) + if gathered_latent.ndim != 2 or gathered_latent.shape[1] != MLA_LATENT_DIM: + raise ValueError("gathered_latent must have shape [capacity, 512]") + if gathered_rope.ndim != 2 or gathered_rope.shape[1] != MLA_ROPE_DIM: + raise ValueError("gathered_rope must have shape [capacity, 64]") + if gathered_latent.shape[0] != gathered_rope.shape[0]: + raise ValueError("gathered_latent and gathered_rope capacities must match") + + +def validate_gather_metadata( + active_slots: torch.Tensor, + request_indices: torch.Tensor, + context_lens: torch.Tensor, + packed_start_locs: torch.Tensor, + *, + cache_slot_count: int, + output_capacity: int, + max_context_len: int, +) -> None: + """Synchronously validate one gather description before per-layer reuse.""" + + if max_context_len < 0 or max_context_len > active_slots.shape[1]: + raise ValueError( + "max_context_len must be within active_slots capacity: " + f"{max_context_len} > {active_slots.shape[1]}" + ) + request_rows = request_indices.tolist() + lengths = context_lens.tolist() + packed_starts = packed_start_locs.tolist() + intervals: list[tuple[int, int]] = [] + for batch_index, (request_row, length, packed_start) in enumerate( + zip(request_rows, lengths, packed_starts) + ): + if length < 0 or length > max_context_len: + raise ValueError( + f"context_lens[{batch_index}]={length} is outside " + f"[0, {max_context_len}]" + ) + if request_row < 0: + if length != 0: + raise ValueError( + "padded request rows must have zero context length" + ) + continue + if request_row >= active_slots.shape[0]: + raise ValueError( + f"request_indices[{batch_index}]={request_row} is outside " + f"[0, {active_slots.shape[0]})" + ) + if packed_start < 0 or packed_start + length > output_capacity: + raise ValueError( + f"packed output for batch {batch_index} exceeds capacity " + f"{output_capacity}" + ) + if length == 0: + continue + slots = active_slots[request_row, :length] + if bool(torch.any(slots < 0).item()) or bool( + torch.any(slots >= cache_slot_count).item() + ): + raise ValueError( + f"active_slots row {request_row} contains an invalid slot" + ) + if slots.unique().numel() != slots.numel(): + raise ValueError( + f"active_slots row {request_row} contains duplicate slots" + ) + intervals.append((packed_start, packed_start + length)) + + intervals.sort() + for previous, current in zip(intervals, intervals[1:]): + if current[0] < previous[1]: + raise ValueError("packed gather output ranges overlap") + + +@torch.no_grad() +def gather_latent_history( + latent_cache: torch.Tensor, + rope_cache: torch.Tensor, + active_slots: torch.Tensor, + request_indices: torch.Tensor, + context_lens: torch.Tensor, + packed_start_locs: torch.Tensor, + gathered_latent: torch.Tensor, + gathered_rope: torch.Tensor, + *, + max_context_len: int, + validate_metadata: bool = True, +) -> None: + """Gather complete ragged MLA history into caller-owned packed buffers.""" + + _validate_gather_tensors( + latent_cache, + rope_cache, + active_slots, + request_indices, + context_lens, + packed_start_locs, + gathered_latent, + gathered_rope, + ) + output_capacity = gathered_latent.shape[0] + if validate_metadata: + validate_gather_metadata( + active_slots, + request_indices, + context_lens, + packed_start_locs, + cache_slot_count=latent_cache.shape[0], + output_capacity=output_capacity, + max_context_len=max_context_len, + ) + if max_context_len == 0: + return + + block_seq = 64 + batch_size = context_lens.numel() + _gather_latent_kernel[ + (batch_size, triton.cdiv(max_context_len, block_seq)) + ]( + latent_cache, + rope_cache, + active_slots, + request_indices, + context_lens, + packed_start_locs, + gathered_latent, + gathered_rope, + *latent_cache.stride(), + *rope_cache.stride(), + *active_slots.stride(), + request_indices.stride(0), + context_lens.stride(0), + packed_start_locs.stride(0), + *gathered_latent.stride(), + *gathered_rope.stride(), + cache_slot_count=latent_cache.shape[0], + output_capacity=output_capacity, + LATENT_DIM=MLA_LATENT_DIM, + ROPE_DIM=MLA_ROPE_DIM, + BLOCK_SEQ=block_seq, + num_warps=8, + num_stages=1, + ) diff --git a/tests/test_mla_kernels.py b/tests/test_mla_kernels.py new file mode 100644 index 00000000..47c5b7ae --- /dev/null +++ b/tests/test_mla_kernels.py @@ -0,0 +1,580 @@ +from __future__ import annotations + +import ast +from dataclasses import replace +from pathlib import Path + +import pytest +import torch + +from sparsevllm.triton_kernel.mla import ( + DEFAULT_GLM_MLA_DECODE_CONFIG, + GLM_MLA_SOFTMAX_SCALE, + MlaDecodeWorkspace, + allocate_mla_decode_workspace, + copy_latent_to_cache, + decode_stage1, + decode_stage2, + gather_latent_history, + prepare_mla_decode_schedule, + run_mla_decode, +) + + +CUDA_REQUIRED = pytest.mark.skipif( + not torch.cuda.is_available(), + reason="CUDA is required for MLA Triton tests", +) +DECODE_CONTEXTS = (1, 31, 32, 33, 127, 128, 129, 255, 256, 257, 1024, 4096) + + +def _torch_mla_decode( + q_latent: torch.Tensor, + q_rope: torch.Tensor, + latent_cache: torch.Tensor, + rope_cache: torch.Tensor, + active_slots: torch.Tensor, + request_indices: torch.Tensor, + context_lens: torch.Tensor, +) -> torch.Tensor: + rows = [] + for batch_index in range(q_latent.shape[0]): + context_len = int(context_lens[batch_index].item()) + if context_len == 0: + rows.append(torch.zeros_like(q_latent[batch_index])) + continue + request_row = int(request_indices[batch_index].item()) + slots = active_slots[request_row, :context_len].long() + keys_latent = latent_cache[slots, 0].float() + keys_rope = rope_cache[slots, 0].float() + logits = torch.matmul(q_latent[batch_index].float(), keys_latent.T) + logits += torch.matmul(q_rope[batch_index].float(), keys_rope.T) + probabilities = torch.softmax( + logits * GLM_MLA_SOFTMAX_SCALE, + dim=-1, + ) + rows.append(torch.matmul(probabilities, keys_latent).to(torch.bfloat16)) + return torch.stack(rows) + + +def _make_decode_case( + batch_size: int, + head_count: int, + max_context_len: int, +) -> tuple[torch.Tensor, ...]: + device = torch.device("cuda") + lengths = [ + max(1, max_context_len - ((batch_index * 17) % max(1, max_context_len // 3 + 1))) + for batch_index in range(batch_size) + ] + request_rows = list(reversed(range(batch_size))) + slot_count = sum(lengths) + 31 + physical_slots = torch.randperm(slot_count, dtype=torch.int32) + active_slots_cpu = torch.full( + (batch_size, max_context_len), + -1, + dtype=torch.int32, + ) + cursor = 0 + for batch_index, (request_row, length) in enumerate( + zip(request_rows, lengths) + ): + del batch_index + active_slots_cpu[request_row, :length] = physical_slots[ + cursor : cursor + length + ] + cursor += length + + q_latent = torch.randn( + (batch_size, head_count, 512), + dtype=torch.bfloat16, + device=device, + ) + q_rope = torch.randn( + (batch_size, head_count, 64), + dtype=torch.bfloat16, + device=device, + ) + latent_cache = torch.randn( + (slot_count, 1, 512), + dtype=torch.bfloat16, + device=device, + ) + rope_cache = torch.randn( + (slot_count, 1, 64), + dtype=torch.bfloat16, + device=device, + ) + active_slots = active_slots_cpu.to(device) + request_indices = torch.tensor( + request_rows, + dtype=torch.int32, + device=device, + ) + context_lens = torch.tensor(lengths, dtype=torch.int32, device=device) + return ( + q_latent, + q_rope, + latent_cache, + rope_cache, + active_slots, + request_indices, + context_lens, + ) + + +def test_vendor_python_files_do_not_import_lightllm() -> None: + kernel_dir = ( + Path(__file__).parents[1] + / "src" + / "sparsevllm" + / "triton_kernel" + / "mla" + ) + for path in kernel_dir.glob("*.py"): + tree = ast.parse(path.read_text(), filename=str(path)) + for node in ast.walk(tree): + if isinstance(node, ast.Import): + imported = [alias.name for alias in node.names] + elif isinstance(node, ast.ImportFrom): + imported = [node.module or ""] + else: + continue + assert all(not name.startswith("lightllm") for name in imported) + + +@CUDA_REQUIRED +def test_copy_latent_skips_padding_and_supports_strides() -> None: + torch.manual_seed(3) + latent_source = torch.randn( + (5, 1, 1024), + dtype=torch.bfloat16, + device="cuda", + ) + rope_source = torch.randn( + (5, 1, 128), + dtype=torch.bfloat16, + device="cuda", + ) + latent = latent_source[..., ::2] + rope = rope_source[..., ::2] + assert not latent.is_contiguous() + assert not rope.is_contiguous() + slots = torch.tensor([2, -1, 5, 0, -1], dtype=torch.int32, device="cuda") + latent_cache = torch.full( + (8, 1, 512), + 7.0, + dtype=torch.bfloat16, + device="cuda", + ) + rope_cache = torch.full( + (8, 1, 64), + 9.0, + dtype=torch.bfloat16, + device="cuda", + ) + + copy_latent_to_cache(latent, rope, slots, latent_cache, rope_cache) + torch.cuda.synchronize() + + torch.testing.assert_close(latent_cache[2], latent[0]) + torch.testing.assert_close(latent_cache[5], latent[2]) + torch.testing.assert_close(latent_cache[0], latent[3]) + torch.testing.assert_close( + latent_cache[1], + torch.full_like(latent_cache[1], 7.0), + ) + torch.testing.assert_close(rope_cache[2], rope[0]) + torch.testing.assert_close(rope_cache[5], rope[2]) + torch.testing.assert_close(rope_cache[0], rope[3]) + torch.testing.assert_close( + rope_cache[1], + torch.full_like(rope_cache[1], 9.0), + ) + + +@CUDA_REQUIRED +def test_copy_latent_rejects_duplicate_and_out_of_range_slots() -> None: + latent = torch.zeros((2, 1, 512), dtype=torch.bfloat16, device="cuda") + rope = torch.zeros((2, 1, 64), dtype=torch.bfloat16, device="cuda") + latent_cache = torch.zeros((4, 1, 512), dtype=torch.bfloat16, device="cuda") + rope_cache = torch.zeros((4, 1, 64), dtype=torch.bfloat16, device="cuda") + + with pytest.raises(ValueError, match="duplicate"): + copy_latent_to_cache( + latent, + rope, + torch.tensor([1, 1], dtype=torch.int32, device="cuda"), + latent_cache, + rope_cache, + ) + with pytest.raises(ValueError, match="outside"): + copy_latent_to_cache( + latent, + rope, + torch.tensor([0, 4], dtype=torch.int32, device="cuda"), + latent_cache, + rope_cache, + ) + + +@CUDA_REQUIRED +def test_gather_latent_full_ragged_history_with_padded_row() -> None: + torch.manual_seed(5) + latent_cache = torch.randn( + (24, 1, 512), + dtype=torch.bfloat16, + device="cuda", + ) + rope_cache = torch.randn( + (24, 1, 64), + dtype=torch.bfloat16, + device="cuda", + ) + active_storage = torch.full((3, 14), -1, dtype=torch.int32, device="cuda") + active_slots = active_storage[:, ::2] + assert not active_slots.is_contiguous() + row_two = torch.tensor([7, 1, 13, 4, 18], dtype=torch.int32, device="cuda") + row_zero = torch.tensor([6, 21, 2], dtype=torch.int32, device="cuda") + active_slots[2, :5] = row_two + active_slots[0, :3] = row_zero + request_indices = torch.tensor([2, 0, -1], dtype=torch.int32, device="cuda") + context_lens = torch.tensor([5, 3, 0], dtype=torch.int32, device="cuda") + packed_starts = torch.tensor([0, 5, 8], dtype=torch.int32, device="cuda") + gathered_latent = torch.full( + (8, 512), + float("nan"), + dtype=torch.bfloat16, + device="cuda", + ) + gathered_rope = torch.full( + (8, 64), + float("nan"), + dtype=torch.bfloat16, + device="cuda", + ) + + gather_latent_history( + latent_cache, + rope_cache, + active_slots, + request_indices, + context_lens, + packed_starts, + gathered_latent, + gathered_rope, + max_context_len=7, + ) + torch.cuda.synchronize() + + expected_slots = torch.cat((row_two, row_zero)).long() + torch.testing.assert_close(gathered_latent, latent_cache[expected_slots, 0]) + torch.testing.assert_close(gathered_rope, rope_cache[expected_slots, 0]) + + +@CUDA_REQUIRED +def test_gather_latent_rejects_duplicate_source_positions() -> None: + latent_cache = torch.zeros((4, 1, 512), dtype=torch.bfloat16, device="cuda") + rope_cache = torch.zeros((4, 1, 64), dtype=torch.bfloat16, device="cuda") + active_slots = torch.tensor([[1, 1]], dtype=torch.int32, device="cuda") + request_indices = torch.tensor([0], dtype=torch.int32, device="cuda") + context_lens = torch.tensor([2], dtype=torch.int32, device="cuda") + packed_starts = torch.tensor([0], dtype=torch.int32, device="cuda") + gathered_latent = torch.empty((2, 512), dtype=torch.bfloat16, device="cuda") + gathered_rope = torch.empty((2, 64), dtype=torch.bfloat16, device="cuda") + + with pytest.raises(ValueError, match="duplicate"): + gather_latent_history( + latent_cache, + rope_cache, + active_slots, + request_indices, + context_lens, + packed_starts, + gathered_latent, + gathered_rope, + max_context_len=2, + ) + + +@pytest.mark.parametrize("batch_size", (1, 2, 8)) +@pytest.mark.parametrize("head_count", (5, 10, 20)) +@CUDA_REQUIRED +def test_mla_decode_matches_torch_matrix( + batch_size: int, + head_count: int, +) -> None: + torch.manual_seed(17 + batch_size * 10 + head_count) + for max_context_len in DECODE_CONTEXTS: + case = _make_decode_case(batch_size, head_count, max_context_len) + q_latent, q_rope, latent_cache, rope_cache = case[:4] + active_slots, request_indices, context_lens = case[4:] + output = torch.empty_like(q_latent) + workspace = allocate_mla_decode_workspace( + batch_size=batch_size, + head_count=head_count, + device=q_latent.device, + ) + + run_mla_decode( + q_latent, + q_rope, + latent_cache, + rope_cache, + active_slots, + request_indices, + context_lens, + output, + workspace, + ) + torch.cuda.synchronize() + expected = _torch_mla_decode(*case) + torch.testing.assert_close( + output.float(), + expected.float(), + rtol=3e-2, + atol=3e-2, + msg=lambda message: ( + f"batch={batch_size}, heads={head_count}, " + f"context={max_context_len}: {message}" + ), + ) + + +@CUDA_REQUIRED +def test_mla_decode_zeroes_padded_rows() -> None: + torch.manual_seed(19) + batch_size = 3 + head_count = 5 + q_latent = torch.randn( + (batch_size, head_count, 512), + dtype=torch.bfloat16, + device="cuda", + ) + q_rope = torch.randn( + (batch_size, head_count, 64), + dtype=torch.bfloat16, + device="cuda", + ) + latent_cache = torch.randn((12, 1, 512), dtype=torch.bfloat16, device="cuda") + rope_cache = torch.randn((12, 1, 64), dtype=torch.bfloat16, device="cuda") + active_slots = torch.full((2, 5), -1, dtype=torch.int32, device="cuda") + active_slots[0, :5] = torch.tensor([1, 3, 5, 7, 9], device="cuda") + active_slots[1, :3] = torch.tensor([2, 4, 6], device="cuda") + request_indices = torch.tensor([0, -1, 1], dtype=torch.int32, device="cuda") + context_lens = torch.tensor([5, 0, 3], dtype=torch.int32, device="cuda") + output = torch.empty_like(q_latent) + workspace = allocate_mla_decode_workspace( + batch_size=batch_size, + head_count=head_count, + device="cuda", + ) + + run_mla_decode( + q_latent, + q_rope, + latent_cache, + rope_cache, + active_slots, + request_indices, + context_lens, + output, + workspace, + ) + torch.cuda.synchronize() + + expected = _torch_mla_decode( + q_latent, + q_rope, + latent_cache, + rope_cache, + active_slots, + request_indices, + context_lens, + ) + torch.testing.assert_close(output.float(), expected.float(), rtol=3e-2, atol=3e-2) + torch.testing.assert_close(output[1], torch.zeros_like(output[1])) + + +@CUDA_REQUIRED +def test_mla_decode_rejects_duplicate_active_slots() -> None: + q_latent = torch.zeros((1, 5, 512), dtype=torch.bfloat16, device="cuda") + q_rope = torch.zeros((1, 5, 64), dtype=torch.bfloat16, device="cuda") + latent_cache = torch.zeros((4, 1, 512), dtype=torch.bfloat16, device="cuda") + rope_cache = torch.zeros((4, 1, 64), dtype=torch.bfloat16, device="cuda") + active_slots = torch.tensor([[1, 1]], dtype=torch.int32, device="cuda") + request_indices = torch.tensor([0], dtype=torch.int32, device="cuda") + context_lens = torch.tensor([2], dtype=torch.int32, device="cuda") + output = torch.empty_like(q_latent) + workspace = allocate_mla_decode_workspace( + batch_size=1, + head_count=5, + device="cuda", + ) + + with pytest.raises(ValueError, match="duplicate"): + run_mla_decode( + q_latent, + q_rope, + latent_cache, + rope_cache, + active_slots, + request_indices, + context_lens, + output, + workspace, + ) + + +@CUDA_REQUIRED +def test_decode_stage1_matches_per_block_oracle() -> None: + torch.manual_seed(23) + case = _make_decode_case(batch_size=2, head_count=5, max_context_len=129) + q_latent, q_rope, latent_cache, rope_cache = case[:4] + active_slots, request_indices, context_lens = case[4:] + config = replace( + DEFAULT_GLM_MLA_DECODE_CONFIG, + program_count=8, + blocks_per_program=2, + ) + workspace = allocate_mla_decode_workspace( + batch_size=2, + head_count=5, + device="cuda", + config=config, + ) + prepare_mla_decode_schedule(context_lens, workspace, config=config) + decode_stage1( + q_latent, + q_rope, + latent_cache, + rope_cache, + active_slots, + request_indices, + context_lens, + workspace.block_size, + workspace.mid_output, + workspace.mid_logsumexp, + softmax_scale=GLM_MLA_SOFTMAX_SCALE, + program_count=config.program_count, + block_q_heads=config.block_q_heads, + block_n=config.block_n, + pipeline_stages=config.stage1_pipeline_stages, + num_warps=config.stage1_num_warps, + ) + torch.cuda.synchronize() + + block_size = int(workspace.block_size.item()) + starts = workspace.batch_start_indices.cpu().tolist() + lengths = context_lens.cpu().tolist() + for batch_index, (start, length) in enumerate(zip(starts, lengths)): + request_row = int(request_indices[batch_index].item()) + all_slots = active_slots[request_row, :length].long() + for block_index, token_start in enumerate(range(0, length, block_size)): + slots = all_slots[token_start : token_start + block_size] + keys_latent = latent_cache[slots, 0].float() + keys_rope = rope_cache[slots, 0].float() + logits = torch.matmul(q_latent[batch_index].float(), keys_latent.T) + logits += torch.matmul(q_rope[batch_index].float(), keys_rope.T) + logits *= GLM_MLA_SOFTMAX_SCALE + probabilities = torch.softmax(logits, dim=-1) + expected_output = torch.matmul( + probabilities.to(torch.bfloat16).float(), + keys_latent, + ) + expected_lse = torch.logsumexp(logits, dim=-1) + output_index = start + block_index + torch.testing.assert_close( + workspace.mid_output[:5, output_index], + expected_output, + rtol=3e-2, + atol=3e-2, + ) + torch.testing.assert_close( + workspace.mid_logsumexp[:5, output_index], + expected_lse, + rtol=2e-2, + atol=2e-2, + ) + + +@CUDA_REQUIRED +def test_decode_stage2_matches_weighted_block_oracle() -> None: + torch.manual_seed(29) + head_count = 5 + context_lens = torch.tensor([33, 65], dtype=torch.int32, device="cuda") + block_size = torch.tensor([32], dtype=torch.int32, device="cuda") + batch_starts = torch.tensor([0, 2], dtype=torch.int32, device="cuda") + mid_output = torch.randn((head_count, 5, 512), dtype=torch.float32, device="cuda") + mid_lse = torch.randn((head_count, 5), dtype=torch.float32, device="cuda") + output = torch.empty((2, head_count, 512), dtype=torch.bfloat16, device="cuda") + + decode_stage2( + block_size, + batch_starts, + context_lens, + mid_output, + mid_lse, + output, + pipeline_stages=2, + num_warps=4, + ) + torch.cuda.synchronize() + + expected_rows = [] + for batch_index, (start, block_count) in enumerate(((0, 2), (2, 3))): + del batch_index + weights = torch.softmax(mid_lse[:, start : start + block_count], dim=-1) + expected_rows.append( + torch.einsum( + "hb,hbd->hd", + weights, + mid_output[:, start : start + block_count], + ).to(torch.bfloat16) + ) + expected = torch.stack(expected_rows) + torch.testing.assert_close(output.float(), expected.float(), rtol=1e-2, atol=1e-2) + + +@CUDA_REQUIRED +def test_decode_rejects_small_workspace_and_non_bf16_input() -> None: + case = _make_decode_case(batch_size=1, head_count=5, max_context_len=33) + q_latent, q_rope, latent_cache, rope_cache = case[:4] + active_slots, request_indices, context_lens = case[4:] + output = torch.empty_like(q_latent) + workspace = allocate_mla_decode_workspace( + batch_size=1, + head_count=5, + device="cuda", + ) + small_workspace = MlaDecodeWorkspace( + block_size=workspace.block_size, + batch_start_indices=workspace.batch_start_indices, + mid_output=workspace.mid_output[:, :-1], + mid_logsumexp=workspace.mid_logsumexp[:, :-1], + ) + + with pytest.raises(ValueError, match="workspace is too small"): + run_mla_decode( + q_latent, + q_rope, + latent_cache, + rope_cache, + active_slots, + request_indices, + context_lens, + output, + small_workspace, + ) + with pytest.raises(TypeError, match="q_latent"): + run_mla_decode( + q_latent.float(), + q_rope, + latent_cache, + rope_cache, + active_slots, + request_indices, + context_lens, + output, + workspace, + ) From eb372514b06e301483da54093804a850b11df3fb Mon Sep 17 00:00:00 2001 From: Codex Date: Tue, 4 Aug 2026 21:56:42 +0800 Subject: [PATCH 02/48] refactor: tag attention payload views --- .../engine/cache_manager/__init__.py | 16 +- src/sparsevllm/engine/cache_manager/base.py | 83 ++++--- .../engine/cache_manager/deltakv_base.py | 44 ++-- .../cache_manager/deltakv_less_memory.py | 69 ++++-- src/sparsevllm/engine/cache_manager/h2o.py | 23 +- src/sparsevllm/engine/cache_manager/rkv.py | 11 +- src/sparsevllm/engine/cache_manager/snapkv.py | 43 ++-- src/sparsevllm/layers/attention.py | 35 ++- src/sparsevllm/layers/attention_backend.py | 233 ++++++++++-------- tests/test_fake_attention_backend.py | 103 ++++++-- tests/test_h2o_cache_manager.py | 34 ++- tests/test_prefill_schedule_policy.py | 25 +- tests/test_qwen35_hd256_decode_routing.py | 21 +- tests/test_rkv_skipkv_methods.py | 20 +- 14 files changed, 500 insertions(+), 260 deletions(-) diff --git a/src/sparsevllm/engine/cache_manager/__init__.py b/src/sparsevllm/engine/cache_manager/__init__.py index 79a1a02b..132dd24f 100644 --- a/src/sparsevllm/engine/cache_manager/__init__.py +++ b/src/sparsevllm/engine/cache_manager/__init__.py @@ -1,11 +1,25 @@ from __future__ import annotations -from .base import CacheManager, DecodeComputeView, LayerBatchStates, PrefillComputeView, SparseSelection +from .base import ( + AttentionPayload, + AttentionViewMeta, + CacheManager, + DecodeComputeView, + ExplicitKVPayload, + LayerBatchStates, + MlaLatentPayload, + PrefillComputeView, + SparseSelection, +) __all__ = [ + "AttentionPayload", + "AttentionViewMeta", "CacheManager", "DecodeComputeView", + "ExplicitKVPayload", "LayerBatchStates", + "MlaLatentPayload", "PrefillComputeView", "SparseSelection", "StandardCacheManager", diff --git a/src/sparsevllm/engine/cache_manager/base.py b/src/sparsevllm/engine/cache_manager/base.py index d4fffd6c..5fab60f6 100644 --- a/src/sparsevllm/engine/cache_manager/base.py +++ b/src/sparsevllm/engine/cache_manager/base.py @@ -141,34 +141,53 @@ class SparseSelection: release_temp_slots: bool = False -@dataclass -class DecodeComputeView: - """Physical KV/view tensors consumed by decode attention kernels.""" +@dataclass(frozen=True) +class AttentionViewMeta: + """Logical request/slot coordinates shared by attention payloads.""" - k_cache: torch.Tensor - v_cache: torch.Tensor active_slots: torch.Tensor req_indices: torch.Tensor context_lens: torch.Tensor - attn_score: torch.Tensor | None = None max_context_len: int | None = None + attn_score: torch.Tensor | None = None temp_slots: torch.Tensor | None = None + + +@dataclass(frozen=True) +class ExplicitKVPayload: + """Materialized key/value tensors consumed by ordinary attention.""" + + k_cache: torch.Tensor + v_cache: torch.Tensor backend: str = "dense" metadata: dict[str, Any] | None = None -@dataclass +@dataclass(frozen=True) +class MlaLatentPayload: + """Latent and RoPE caches consumed by MLA attention providers.""" + + latent_cache: torch.Tensor + rope_cache: torch.Tensor + + +AttentionPayload = ExplicitKVPayload | MlaLatentPayload + + +@dataclass(frozen=True) +class DecodeComputeView: + """Decode metadata paired with exactly one physical payload layout.""" + + meta: AttentionViewMeta + payload: AttentionPayload + + +@dataclass(frozen=True) class PrefillComputeView: - """Physical KV/view tensors consumed by prefill attention kernels.""" + """Prefill metadata paired with exactly one physical payload layout.""" - k_cache: torch.Tensor - v_cache: torch.Tensor - active_slots: torch.Tensor - req_indices: torch.Tensor - context_lens: torch.Tensor - attn_score: torch.Tensor | None = None - max_context_len: int | None = None - temp_slots: torch.Tensor | None = None + meta: AttentionViewMeta + payload: AttentionPayload class CacheManager(ABC): @@ -675,14 +694,15 @@ def build_prefill_compute_view( context_lens, ) return PrefillComputeView( - k_cache=k_cache, - v_cache=v_cache, - active_slots=active_slots, - req_indices=req_indices, - context_lens=context_lens, - attn_score=selection.attn_score, - max_context_len=selection.max_context_len, - temp_slots=temp_slots, + meta=AttentionViewMeta( + active_slots=active_slots, + req_indices=req_indices, + context_lens=context_lens, + attn_score=selection.attn_score, + max_context_len=selection.max_context_len, + temp_slots=temp_slots, + ), + payload=ExplicitKVPayload(k_cache=k_cache, v_cache=v_cache), ) def collect_prefill_attention_score( @@ -921,13 +941,14 @@ def build_decode_compute_view( selection, ) return DecodeComputeView( - k_cache=k_cache, - v_cache=v_cache, - active_slots=active_slots, - req_indices=req_indices, - context_lens=context_lens, - attn_score=selection.attn_score, - max_context_len=selection.max_context_len, + meta=AttentionViewMeta( + active_slots=active_slots, + req_indices=req_indices, + context_lens=context_lens, + attn_score=selection.attn_score, + max_context_len=selection.max_context_len, + ), + payload=ExplicitKVPayload(k_cache=k_cache, v_cache=v_cache), ) def get_decode_block_seq(self, layer_idx: int, default: int) -> int: diff --git a/src/sparsevllm/engine/cache_manager/deltakv_base.py b/src/sparsevllm/engine/cache_manager/deltakv_base.py index 3d33acac..9687453e 100644 --- a/src/sparsevllm/engine/cache_manager/deltakv_base.py +++ b/src/sparsevllm/engine/cache_manager/deltakv_base.py @@ -24,7 +24,15 @@ from sparsevllm.layers.rotary_embedding import get_rope, apply_rotary_emb from sparsevllm.platforms import device_runtime -from .base import CacheManager, DecodeComputeView, LayerBatchStates, PrefillComputeView, SparseSelection +from .base import ( + AttentionViewMeta, + CacheManager, + DecodeComputeView, + ExplicitKVPayload, + LayerBatchStates, + PrefillComputeView, + SparseSelection, +) from .raw_kv_offload import RawKVOffloadBuffer @@ -953,14 +961,15 @@ def build_prefill_compute_view( context_lens, ) return PrefillComputeView( - k_cache=k_cache, - v_cache=v_cache, - active_slots=active_slots, - req_indices=req_indices, - context_lens=context_lens, - attn_score=selection.attn_score, - max_context_len=selection.max_context_len, - temp_slots=temp_slots, + meta=AttentionViewMeta( + active_slots=active_slots, + req_indices=req_indices, + context_lens=context_lens, + attn_score=selection.attn_score, + max_context_len=selection.max_context_len, + temp_slots=temp_slots, + ), + payload=ExplicitKVPayload(k_cache=k_cache, v_cache=v_cache), ) def build_decode_compute_view( @@ -997,14 +1006,15 @@ def build_decode_compute_view( selection, ) return DecodeComputeView( - k_cache=k_cache, - v_cache=v_cache, - active_slots=active_slots, - req_indices=req_indices, - context_lens=context_lens, - attn_score=selection.attn_score, - max_context_len=selection.max_context_len, - temp_slots=temp_slots, + meta=AttentionViewMeta( + active_slots=active_slots, + req_indices=req_indices, + context_lens=context_lens, + attn_score=selection.attn_score, + max_context_len=selection.max_context_len, + temp_slots=temp_slots, + ), + payload=ExplicitKVPayload(k_cache=k_cache, v_cache=v_cache), ) def has_prefill_staging_view(self, layer_idx: int) -> bool: diff --git a/src/sparsevllm/engine/cache_manager/deltakv_less_memory.py b/src/sparsevllm/engine/cache_manager/deltakv_less_memory.py index 895e8c7c..22dd57a4 100644 --- a/src/sparsevllm/engine/cache_manager/deltakv_less_memory.py +++ b/src/sparsevllm/engine/cache_manager/deltakv_less_memory.py @@ -2,6 +2,7 @@ import math import os +from dataclasses import replace import torch @@ -20,7 +21,12 @@ from sparsevllm.utils.log import logger from sparsevllm.utils.profiler import profiler -from .base import DecodeComputeView, SparseSelection +from .base import ( + AttentionViewMeta, + DecodeComputeView, + ExplicitKVPayload, + SparseSelection, +) from .deltakv_base import DeltaKVCacheTritonManagerV4 @@ -2947,28 +2953,32 @@ def build_decode_compute_view( if self.full_layer_kivi_key_packed is None or self.full_layer_kivi_value_packed is None: raise RuntimeError("Full-layer KIVI decode was requested before KIVI storage was initialized.") return DecodeComputeView( - k_cache=self.full_kv_cache[0, l_idx], - v_cache=self.full_kv_cache[1, l_idx], - active_slots=self.full_layer_slots_map, - req_indices=selection.req_indices, - context_lens=selection.context_lens, - attn_score=selection.attn_score, - max_context_len=selection.max_context_len, - backend="full_layer_kivi", - metadata={ - "kivi_block_slots_map": self.full_layer_kivi_block_slots_map, - "kivi_block_start_pos": self.full_layer_kivi_block_start_pos, - "key_packed": self.full_layer_kivi_key_packed[l_idx], - "key_scales": self.full_layer_kivi_key_scales[l_idx], - "key_mins": self.full_layer_kivi_key_mins[l_idx], - "value_packed": self.full_layer_kivi_value_packed[l_idx], - "value_scales": self.full_layer_kivi_value_scales[l_idx], - "value_mins": self.full_layer_kivi_value_mins[l_idx], - "group_size": self._full_layer_kivi_group_size(), - "block_n": int(getattr(self.config, "full_layer_kivi_decode_block_n", 16) or 16), - "num_warps": int(getattr(self.config, "full_layer_kivi_decode_num_warps", 2) or 2), - "num_stages": int(getattr(self.config, "full_layer_kivi_decode_num_stages", 3) or 3), - }, + meta=AttentionViewMeta( + active_slots=self.full_layer_slots_map, + req_indices=selection.req_indices, + context_lens=selection.context_lens, + attn_score=selection.attn_score, + max_context_len=selection.max_context_len, + ), + payload=ExplicitKVPayload( + k_cache=self.full_kv_cache[0, l_idx], + v_cache=self.full_kv_cache[1, l_idx], + backend="full_layer_kivi", + metadata={ + "kivi_block_slots_map": self.full_layer_kivi_block_slots_map, + "kivi_block_start_pos": self.full_layer_kivi_block_start_pos, + "key_packed": self.full_layer_kivi_key_packed[l_idx], + "key_scales": self.full_layer_kivi_key_scales[l_idx], + "key_mins": self.full_layer_kivi_key_mins[l_idx], + "value_packed": self.full_layer_kivi_value_packed[l_idx], + "value_scales": self.full_layer_kivi_value_scales[l_idx], + "value_mins": self.full_layer_kivi_value_mins[l_idx], + "group_size": self._full_layer_kivi_group_size(), + "block_n": int(getattr(self.config, "full_layer_kivi_decode_block_n", 16) or 16), + "num_warps": int(getattr(self.config, "full_layer_kivi_decode_num_warps", 2) or 2), + "num_stages": int(getattr(self.config, "full_layer_kivi_decode_num_stages", 3) or 3), + }, + ), ) view = super().build_decode_compute_view( layer_idx, @@ -2982,7 +2992,18 @@ def build_decode_compute_view( and not self.has_prefill_staging_view(layer_idx) and getattr(self.config, "deltakv_sparse_decode_backend", "custom") == "fa2" ): - view.backend = "flash_attn_contiguous" + if not isinstance(view.payload, ExplicitKVPayload): + raise TypeError( + "DeltaKV FA2 decode requires ExplicitKVPayload, got " + f"{type(view.payload).__name__}." + ) + view = replace( + view, + payload=replace( + view.payload, + backend="flash_attn_contiguous", + ), + ) return view def get_layer_compute_tensors(self, layer_idx: int, selection: SparseSelection | None = None): diff --git a/src/sparsevllm/engine/cache_manager/h2o.py b/src/sparsevllm/engine/cache_manager/h2o.py index 4f199069..e935f249 100644 --- a/src/sparsevllm/engine/cache_manager/h2o.py +++ b/src/sparsevllm/engine/cache_manager/h2o.py @@ -11,7 +11,7 @@ from sparsevllm.utils.context import get_context from sparsevllm.utils.profiler import profiler -from .base import PrefillComputeView +from .base import ExplicitKVPayload, PrefillComputeView from .snapkv import SnapKVCacheManager @@ -754,6 +754,13 @@ def collect_prefill_attention_score( ranges = self.prefill_score_ranges(layer_idx, seqs) if not ranges: return None + if not isinstance(view.payload, ExplicitKVPayload): + raise TypeError( + "H2O prefill scoring requires ExplicitKVPayload, got " + f"{type(view.payload).__name__}." + ) + meta = view.meta + payload = view.payload score_starts = torch.tensor( [item[3] for item in ranges], dtype=torch.int32, device=q.device @@ -765,31 +772,31 @@ def collect_prefill_attention_score( [item[4] for item in ranges], dtype=torch.int32, device=q.device ) physical_coordinates_match = ( - view.context_lens[: len(seqs)] == physical_context_lens + meta.context_lens[: len(seqs)] == physical_context_lens ).all() if physical_coordinates_match.is_cuda: torch._assert_async(physical_coordinates_match) elif not bool(physical_coordinates_match.item()): raise RuntimeError( "H2O prefill score view is not in compressed physical coordinates: " - f"layer={layer_idx} view={view.context_lens.tolist()} " + f"layer={layer_idx} view={meta.context_lens.tolist()} " f"physical={physical_context_lens.tolist()}." ) max_context_len = max(item[4] for item in ranges) step_score = torch.zeros( (len(seqs), max_context_len), dtype=torch.float32, device=q.device ) - prompt_cache_lens = view.context_lens - chunk_lens + prompt_cache_lens = meta.context_lens - chunk_lens prefill_score_fwd( q, - view.k_cache, + payload.k_cache, step_score, - view.req_indices, + meta.req_indices, b_start_loc, - view.context_lens, + meta.context_lens, prompt_cache_lens, max(int(seq.current_chunk_size) for seq in seqs), - view.active_slots, + meta.active_slots, score_starts, score_ends, candidate_start=0, diff --git a/src/sparsevllm/engine/cache_manager/rkv.py b/src/sparsevllm/engine/cache_manager/rkv.py index ac249dcd..378496ed 100644 --- a/src/sparsevllm/engine/cache_manager/rkv.py +++ b/src/sparsevllm/engine/cache_manager/rkv.py @@ -264,10 +264,11 @@ def record_prefill_query( layer_idx = int(layer_idx) cache, positions_cache = self._rkv_layer_query_cache(layer_idx) + meta = view.meta if not bool(getattr(self, "_rkv_vectorized_prefill_query_cache", True)): - batch = int(view.context_lens.numel()) + batch = int(meta.context_lens.numel()) for b_idx in range(batch): - context_len = int(view.context_lens[b_idx].item()) + context_len = int(meta.context_lens[b_idx].item()) chunk_len = int(chunk_lens[b_idx].item()) if context_len <= 0 or chunk_len <= 0: continue @@ -277,7 +278,7 @@ def record_prefill_query( if record_len <= 0: continue - row_idx = int(view.req_indices[b_idx].item()) + row_idx = int(meta.req_indices[b_idx].item()) q_start = int(b_start_loc[b_idx].item()) + (record_start - chunk_start) token_positions = torch.arange( record_start, @@ -290,9 +291,9 @@ def record_prefill_query( positions_cache[row_idx, cols] = token_positions.to(torch.int32) return None - context_lens = view.context_lens.to(device=q.device, dtype=torch.long) + context_lens = meta.context_lens.to(device=q.device, dtype=torch.long) chunk_lens = chunk_lens.to(device=q.device, dtype=torch.long) - req_indices = view.req_indices.to(device=q.device, dtype=torch.long) + req_indices = meta.req_indices.to(device=q.device, dtype=torch.long) b_start_loc = b_start_loc.to(device=q.device, dtype=torch.long) offsets = torch.arange(obs, dtype=torch.long, device=q.device) diff --git a/src/sparsevllm/engine/cache_manager/snapkv.py b/src/sparsevllm/engine/cache_manager/snapkv.py index 13ed9325..baa5e759 100644 --- a/src/sparsevllm/engine/cache_manager/snapkv.py +++ b/src/sparsevllm/engine/cache_manager/snapkv.py @@ -20,7 +20,13 @@ from sparsevllm.utils.log import logger, log_level from sparsevllm.utils.profiler import profiler -from .base import CacheManager, LayerBatchStates, PrefillComputeView, SparseSelection +from .base import ( + CacheManager, + ExplicitKVPayload, + LayerBatchStates, + PrefillComputeView, + SparseSelection, +) from .raw_kv_offload import RawKVOffloadBuffer @@ -969,8 +975,15 @@ def collect_prefill_attention_score( rows = self._prefill_score_rows(layer_idx, seqs) if not rows: return None + if not isinstance(view.payload, ExplicitKVPayload): + raise TypeError( + "SnapKV prefill scoring requires ExplicitKVPayload, got " + f"{type(view.payload).__name__}." + ) + meta = view.meta + payload = view.payload - b_prompt_cache_len = view.context_lens - chunk_lens + b_prompt_cache_len = meta.context_lens - chunk_lens max_query_len = max(int(seq.current_chunk_size) for seq in seqs) if len(rows) == 1: b_idx, seq, score_start, score_end = rows[0] @@ -978,19 +991,19 @@ def collect_prefill_attention_score( acc = self._get_prefill_attention_score_accumulator( layer_idx, seq, - prompt_len=int(view.context_lens[b_idx].item()), + prompt_len=int(meta.context_lens[b_idx].item()), device=q.device, ) prefill_score_fwd( q, - view.k_cache, + payload.k_cache, acc.unsqueeze(0), - view.req_indices[b_idx : b_idx + 1], + meta.req_indices[b_idx : b_idx + 1], b_start_loc[b_idx : b_idx + 1], - view.context_lens[b_idx : b_idx + 1], + meta.context_lens[b_idx : b_idx + 1], b_prompt_cache_len[b_idx : b_idx + 1], query_len, - view.active_slots, + meta.active_slots, *self._prefill_score_bound_tensors( score_start=score_start, score_end=score_end, @@ -1008,8 +1021,8 @@ def collect_prefill_attention_score( score_ends[b_idx] = int(score_end) max_context_len = ( - int(view.max_context_len) - if view.max_context_len is not None + int(meta.max_context_len) + if meta.max_context_len is not None else max(int(seq.num_prefilled_tokens + seq.current_chunk_size) for seq in seqs) ) step_score = torch.zeros( @@ -1019,14 +1032,14 @@ def collect_prefill_attention_score( ) prefill_score_fwd( q, - view.k_cache, + payload.k_cache, step_score, - view.req_indices, + meta.req_indices, b_start_loc, - view.context_lens, + meta.context_lens, b_prompt_cache_len, max_query_len, - view.active_slots, + meta.active_slots, score_starts, score_ends, candidate_start=int(self.config.num_sink_tokens), @@ -1036,10 +1049,10 @@ def collect_prefill_attention_score( acc = self._get_prefill_attention_score_accumulator( layer_idx, seq, - prompt_len=int(view.context_lens[b_idx].item()), + prompt_len=int(meta.context_lens[b_idx].item()), device=q.device, ) - context_len = int(view.context_lens[b_idx].item()) + context_len = int(meta.context_lens[b_idx].item()) acc[:context_len] = torch.maximum(acc[:context_len], step_score[b_idx, :context_len]) return None diff --git a/src/sparsevllm/layers/attention.py b/src/sparsevllm/layers/attention.py index 55731676..4af57348 100644 --- a/src/sparsevllm/layers/attention.py +++ b/src/sparsevllm/layers/attention.py @@ -3,6 +3,7 @@ import torch from torch import nn +from sparsevllm.engine.cache_manager import ExplicitKVPayload from sparsevllm.layers.attention_backend import TritonAttentionBackend from sparsevllm.utils.context import get_context @@ -86,20 +87,26 @@ def forward( v, selection, ) - temp_slots = prefill_view.temp_slots + if not isinstance(prefill_view.payload, ExplicitKVPayload): + raise TypeError( + "Attention prefill requires ExplicitKVPayload, got " + f"{type(prefill_view.payload).__name__}." + ) + prefill_meta = prefill_view.meta + temp_slots = prefill_meta.temp_slots if context.cu_seqlens_q is None or context.cu_seqlens_q.numel() <= 1: return torch.empty_like(q) b_start_loc = context.cu_seqlens_q[:-1] chunk_lens = context.cu_seqlens_q[1:] - context.cu_seqlens_q[:-1] - max_context_len = prefill_view.max_context_len + max_context_len = prefill_meta.max_context_len if max_context_len is not None: max_input_len = int(max_context_len) elif torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): - max_input_len = int(prefill_view.active_slots.shape[1]) + max_input_len = int(prefill_meta.active_slots.shape[1]) else: - max_input_len = prefill_view.context_lens.max().item() + max_input_len = prefill_meta.context_lens.max().item() o = self.attention_backend.run_prefill( q, @@ -135,9 +142,15 @@ def forward( num_heads=self.num_heads, num_kv_heads=self.num_kv_heads, ) - temp_slots = decode_view.temp_slots + if not isinstance(decode_view.payload, ExplicitKVPayload): + raise TypeError( + "Attention decode requires ExplicitKVPayload, got " + f"{type(decode_view.payload).__name__}." + ) + decode_meta = decode_view.meta + temp_slots = decode_meta.temp_slots - max_context_len = decode_view.max_context_len + max_context_len = decode_meta.max_context_len static_cap = getattr(cache_manager, "_decode_static_max_context_len", None) if static_cap is not None: max_context_len = max( @@ -147,13 +160,17 @@ def forward( if max_context_len is None: raise RuntimeError(f"static decode requires max_context_len, got None at layer={layer_idx}") max_len_in_batch = int(max_context_len) - if decode_view.active_slots.dim() == 2: - slot_table_len = int(decode_view.active_slots.shape[1]) + if decode_meta.active_slots.dim() == 2: + slot_table_len = int(decode_meta.active_slots.shape[1]) if ( os.environ.get("SVLLM_DEBUG_DECODE_BOUNDS", "0") == "1" and not (torch.cuda.is_available() and torch.cuda.is_current_stream_capturing()) ): - actual_max_len = int(decode_view.context_lens.max().item()) if decode_view.context_lens.numel() > 0 else 0 + actual_max_len = ( + int(decode_meta.context_lens.max().item()) + if decode_meta.context_lens.numel() > 0 + else 0 + ) if actual_max_len > slot_table_len: raise RuntimeError( "decode context length exceeds active slot table width: " diff --git a/src/sparsevllm/layers/attention_backend.py b/src/sparsevllm/layers/attention_backend.py index 6042cf74..959ee27a 100644 --- a/src/sparsevllm/layers/attention_backend.py +++ b/src/sparsevllm/layers/attention_backend.py @@ -2,7 +2,11 @@ import torch -from sparsevllm.engine.cache_manager import DecodeComputeView, PrefillComputeView +from sparsevllm.engine.cache_manager import ( + DecodeComputeView, + ExplicitKVPayload, + PrefillComputeView, +) from sparsevllm.operators.registry import record_operator_binding from sparsevllm.utils.context import get_context from sparsevllm.triton_kernel.context_flashattention_nopad import context_attention_fwd @@ -92,6 +96,20 @@ def _fill_fake_attention_score(attn_score: torch.Tensor | None) -> None: attn_score.zero_() +def _require_explicit_payload( + view: PrefillComputeView | DecodeComputeView, + *, + operation: str, +) -> ExplicitKVPayload: + payload = view.payload + if not isinstance(payload, ExplicitKVPayload): + raise TypeError( + f"{operation} requires ExplicitKVPayload, got " + f"{type(payload).__name__}." + ) + return payload + + class TritonAttentionBackend: """Thin backend wrapper around the existing Sparse-vLLM Triton attention kernels.""" @@ -109,30 +127,32 @@ def run_prefill( chunk_lens: torch.Tensor, max_input_len: int, ) -> torch.Tensor: - b_seq_len = view.context_lens + payload = _require_explicit_payload(view, operation="Triton prefill") + meta = view.meta + b_seq_len = meta.context_lens if b_seq_len.numel() != chunk_lens.numel(): layer_idx = getattr(get_context(), "now_layer_idx", None) raise RuntimeError( "prefill context_lens/chunk_lens batch mismatch: " f"layer={layer_idx} context_lens_shape={tuple(b_seq_len.shape)} " f"chunk_lens_shape={tuple(chunk_lens.shape)} q_shape={tuple(q.shape)} " - f"req_indices_shape={tuple(view.req_indices.shape)} " - f"active_slots_shape={tuple(view.active_slots.shape)}" + f"req_indices_shape={tuple(meta.req_indices.shape)} " + f"active_slots_shape={tuple(meta.active_slots.shape)}" ) b_prompt_cache_len = b_seq_len - chunk_lens self._debug_check_prefill_bounds(q, view, chunk_lens=chunk_lens) if _fake_prefill_attention_enabled(): real_probe_min_context = _warmup_real_prefill_probe_min_context() probe_context_tokens = ( - int(view.max_context_len) - if view.max_context_len is not None + int(meta.max_context_len) + if meta.max_context_len is not None else int(max_input_len) ) if ( real_probe_min_context is None or probe_context_tokens < real_probe_min_context ): - _fill_fake_attention_score(view.attn_score) + _fill_fake_attention_score(meta.attn_score) return _fake_attention_output(q) log_once( "Warmup real prefill attention probe executing at " @@ -144,16 +164,16 @@ def run_prefill( o = torch.empty_like(q) context_attention_fwd( q, - view.k_cache, - view.v_cache, + payload.k_cache, + payload.v_cache, o, - view.req_indices, + meta.req_indices, b_start_loc, b_seq_len, b_prompt_cache_len, max_input_len, - view.active_slots, - attn_score=view.attn_score, + meta.active_slots, + attn_score=meta.attn_score, ) return o @@ -168,39 +188,41 @@ def _debug_check_prefill_bounds( return if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): return - if view.active_slots.dim() != 2: + payload = _require_explicit_payload(view, operation="Prefill bounds check") + meta = view.meta + if meta.active_slots.dim() != 2: raise RuntimeError( - f"prefill bounds check expects 2D active_slots, got shape={tuple(view.active_slots.shape)}" + f"prefill bounds check expects 2D active_slots, got shape={tuple(meta.active_slots.shape)}" ) - rows = view.req_indices.to(torch.long) + rows = meta.req_indices.to(torch.long) row_min = int(rows.min().item()) if rows.numel() > 0 else 0 row_max = int(rows.max().item()) if rows.numel() > 0 else -1 - if row_min < 0 or row_max >= int(view.active_slots.shape[0]): + if row_min < 0 or row_max >= int(meta.active_slots.shape[0]): raise RuntimeError( "prefill req row index out of bounds: " - f"row_min={row_min} row_max={row_max} num_rows={int(view.active_slots.shape[0])}" + f"row_min={row_min} row_max={row_max} num_rows={int(meta.active_slots.shape[0])}" ) if int(chunk_lens.sum().item()) != int(q.shape[0]): raise RuntimeError( "prefill q/chunk length mismatch: " f"q_tokens={int(q.shape[0])} chunk_tokens={int(chunk_lens.sum().item())}" ) - if bool((view.context_lens < chunk_lens).any().item()): + if bool((meta.context_lens < chunk_lens).any().item()): raise RuntimeError( "prefill context_lens shorter than chunk_lens: " - f"context_lens={view.context_lens.detach().cpu().tolist()} " + f"context_lens={meta.context_lens.detach().cpu().tolist()} " f"chunk_lens={chunk_lens.detach().cpu().tolist()}" ) - visible_len = int(view.context_lens.max().item()) if view.context_lens.numel() > 0 else 0 - if visible_len > int(view.active_slots.shape[1]): + visible_len = int(meta.context_lens.max().item()) if meta.context_lens.numel() > 0 else 0 + if visible_len > int(meta.active_slots.shape[1]): raise RuntimeError( "prefill visible length exceeds active slot table width: " - f"visible_len={visible_len} active_slots_width={int(view.active_slots.shape[1])}" + f"visible_len={visible_len} active_slots_width={int(meta.active_slots.shape[1])}" ) - visible_slots = view.active_slots.index_select(0, rows)[:, :visible_len] + visible_slots = meta.active_slots.index_select(0, rows)[:, :visible_len] pos = torch.arange(visible_len, device=visible_slots.device)[None, :] - valid_pos = pos < view.context_lens[:, None] - slot_cap = int(view.k_cache.shape[0]) + valid_pos = pos < meta.context_lens[:, None] + slot_cap = int(payload.k_cache.shape[0]) bad = ((visible_slots < 0) | (visible_slots >= slot_cap)) & valid_pos if bool(bad.any().item()): layer_idx = getattr(get_context(), "now_layer_idx", None) @@ -212,9 +234,9 @@ def _debug_check_prefill_bounds( raise RuntimeError( "prefill physical slot out of bounds before attention: " f"layer={layer_idx} batch={bad_b} req_row={bad_req_row} pos={bad_pos} " - f"slot={bad_slot} slot_cap={slot_cap} context_len={int(view.context_lens[bad_b].item())} " - f"k_shape={tuple(view.k_cache.shape)} v_shape={tuple(view.v_cache.shape)} " - f"active_slots_shape={tuple(view.active_slots.shape)}" + f"slot={bad_slot} slot_cap={slot_cap} context_len={int(meta.context_lens[bad_b].item())} " + f"k_shape={tuple(payload.k_cache.shape)} v_shape={tuple(payload.v_cache.shape)} " + f"active_slots_shape={tuple(meta.active_slots.shape)}" ) def run_decode( @@ -229,10 +251,12 @@ def run_decode( num_heads: int, num_kv_heads: int, ) -> torch.Tensor: + payload = _require_explicit_payload(view, operation="Triton decode") + meta = view.meta if _fake_decode_attention_enabled(): - _fill_fake_attention_score(view.attn_score) + _fill_fake_attention_score(meta.attn_score) return _fake_attention_output(q) - if view.backend == "full_layer_kivi": + if payload.backend == "full_layer_kivi": self._run_full_layer_kivi_decode_stage1( q, view, @@ -242,30 +266,40 @@ def run_decode( block_seq=block_seq, ) o = torch.empty_like(q) - flash_decode_stage2(mid_o, mid_o_logexpsum, view.context_lens, o, block_seq) + flash_decode_stage2(mid_o, mid_o_logexpsum, meta.context_lens, o, block_seq) return o self._debug_check_decode_bounds(view) - if view.backend == "flash_attn_contiguous": + if payload.backend == "flash_attn_contiguous": from flash_attn import flash_attn_with_kvcache - if view.active_slots.dim() != 2: + if meta.active_slots.dim() != 2: raise RuntimeError("flash_attn_contiguous decode expects a 2D active slot table.") - batch, width = int(view.active_slots.shape[0]), int(view.active_slots.shape[1]) + batch, width = int(meta.active_slots.shape[0]), int(meta.active_slots.shape[1]) expected = batch * width - if int(view.k_cache.shape[0]) < expected or int(view.v_cache.shape[0]) < expected: + if int(payload.k_cache.shape[0]) < expected or int(payload.v_cache.shape[0]) < expected: raise RuntimeError( "flash_attn_contiguous decode got a cache smaller than the materialized active view: " - f"cache={int(view.k_cache.shape[0])}/{int(view.v_cache.shape[0])} expected={expected}." + f"cache={int(payload.k_cache.shape[0])}/{int(payload.v_cache.shape[0])} expected={expected}." ) - k_cache = view.k_cache[:expected].view(batch, width, int(view.k_cache.shape[1]), int(view.k_cache.shape[2])) - v_cache = view.v_cache[:expected].view(batch, width, int(view.v_cache.shape[1]), int(view.v_cache.shape[2])) + k_cache = payload.k_cache[:expected].view( + batch, + width, + int(payload.k_cache.shape[1]), + int(payload.k_cache.shape[2]), + ) + v_cache = payload.v_cache[:expected].view( + batch, + width, + int(payload.v_cache.shape[1]), + int(payload.v_cache.shape[2]), + ) with profiler.record("decode_attention_flash_attn_sparse"): # Decode uses q_len=1 and the materialized KV view contains no future tokens. out = flash_attn_with_kvcache( q.unsqueeze(1), k_cache, v_cache, - cache_seqlens=view.context_lens.to(torch.int32), + cache_seqlens=meta.context_lens.to(torch.int32), causal=False, ) return out.squeeze(1) @@ -273,44 +307,44 @@ def run_decode( profile_kind = "full" if int(max_len_in_batch) > 8192 else "sparse" is_gqa = int(num_heads) > int(num_kv_heads) with profiler.record(f"decode_attention_stage1_{profile_kind}"): - if view.attn_score is not None: + if meta.attn_score is not None: if is_gqa: gqa_flash_decode_stage1_with_score( q, - view.k_cache, - view.v_cache, - view.active_slots, - view.req_indices, - view.context_lens, + payload.k_cache, + payload.v_cache, + meta.active_slots, + meta.req_indices, + meta.context_lens, max_len_in_batch, mid_o, mid_o_logexpsum, - view.attn_score, + meta.attn_score, block_seq, ) else: mha_flash_decode_stage1_with_score( q, - view.k_cache, - view.v_cache, - view.active_slots, - view.req_indices, - view.context_lens, + payload.k_cache, + payload.v_cache, + meta.active_slots, + meta.req_indices, + meta.context_lens, max_len_in_batch, mid_o, mid_o_logexpsum, - view.attn_score, + meta.attn_score, block_seq, ) else: if is_gqa: gqa_flash_decode_stage1( q, - view.k_cache, - view.v_cache, - view.active_slots, - view.req_indices, - view.context_lens, + payload.k_cache, + payload.v_cache, + meta.active_slots, + meta.req_indices, + meta.context_lens, max_len_in_batch, mid_o, mid_o_logexpsum, @@ -319,11 +353,11 @@ def run_decode( else: mha_flash_decode_stage1( q, - view.k_cache, - view.v_cache, - view.active_slots, - view.req_indices, - view.context_lens, + payload.k_cache, + payload.v_cache, + meta.active_slots, + meta.req_indices, + meta.context_lens, max_len_in_batch, mid_o, mid_o_logexpsum, @@ -332,7 +366,7 @@ def run_decode( o = torch.empty_like(q) with profiler.record(f"decode_attention_stage2_{profile_kind}"): - flash_decode_stage2(mid_o, mid_o_logexpsum, view.context_lens, o, block_seq) + flash_decode_stage2(mid_o, mid_o_logexpsum, meta.context_lens, o, block_seq) return o def _run_full_layer_kivi_decode_stage1( @@ -345,66 +379,73 @@ def _run_full_layer_kivi_decode_stage1( max_len_in_batch: int, block_seq: int, ): - meta = view.metadata - if meta is None: + payload = _require_explicit_payload( + view, + operation="Full-layer KIVI decode", + ) + view_meta = view.meta + backend_metadata = payload.metadata + if backend_metadata is None: raise RuntimeError("full_layer_kivi decode view is missing metadata.") from sparsevllm.triton_kernel.deltakv_kernels import full_layer_kivi_flash_decode_stage1 full_layer_kivi_flash_decode_stage1( q=q, - raw_k=view.k_cache, - raw_v=view.v_cache, - raw_slots_map=view.active_slots, - kivi_block_slots_map=meta["kivi_block_slots_map"], - kivi_block_start_pos=meta["kivi_block_start_pos"], - key_packed=meta["key_packed"], - key_scales=meta["key_scales"], - key_mins=meta["key_mins"], - value_packed=meta["value_packed"], - value_scales=meta["value_scales"], - value_mins=meta["value_mins"], - req_indices=view.req_indices, - context_lens=view.context_lens, + raw_k=payload.k_cache, + raw_v=payload.v_cache, + raw_slots_map=view_meta.active_slots, + kivi_block_slots_map=backend_metadata["kivi_block_slots_map"], + kivi_block_start_pos=backend_metadata["kivi_block_start_pos"], + key_packed=backend_metadata["key_packed"], + key_scales=backend_metadata["key_scales"], + key_mins=backend_metadata["key_mins"], + value_packed=backend_metadata["value_packed"], + value_scales=backend_metadata["value_scales"], + value_mins=backend_metadata["value_mins"], + req_indices=view_meta.req_indices, + context_lens=view_meta.context_lens, max_len_in_batch=max_len_in_batch, mid_out=mid_o, mid_out_logsumexp=mid_o_logexpsum, - group_size=int(meta["group_size"]), + group_size=int(backend_metadata["group_size"]), block_seq=block_seq, - block_n=int(meta.get("block_n", 16)), - num_warps=int(meta.get("num_warps", 2)), - num_stages=int(meta.get("num_stages", 3)), - attn_score=view.attn_score, + block_n=int(backend_metadata.get("block_n", 16)), + num_warps=int(backend_metadata.get("num_warps", 2)), + num_stages=int(backend_metadata.get("num_stages", 3)), + attn_score=view_meta.attn_score, ) def _debug_check_decode_bounds(self, view: DecodeComputeView): if os.environ.get("SVLLM_DEBUG_DECODE_BOUNDS", "0") != "1": return - if view.backend not in {"dense", "flash_attn_contiguous"}: + payload = _require_explicit_payload(view, operation="Decode bounds check") + meta = view.meta + if payload.backend not in {"dense", "flash_attn_contiguous"}: return if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): return - if view.active_slots.dim() != 2: + if meta.active_slots.dim() != 2: raise RuntimeError( - f"debug slot bounds check expects 2D active_slots, got shape={tuple(view.active_slots.shape)}" + f"debug slot bounds check expects 2D active_slots, got shape={tuple(meta.active_slots.shape)}" ) - rows = view.req_indices.to(torch.long) + rows = meta.req_indices.to(torch.long) row_min = int(rows.min().item()) if rows.numel() > 0 else 0 row_max = int(rows.max().item()) if rows.numel() > 0 else -1 - if row_min < 0 or row_max >= int(view.active_slots.shape[0]): + if row_min < 0 or row_max >= int(meta.active_slots.shape[0]): raise RuntimeError( "decode req row index out of bounds: " - f"row_min={row_min} row_max={row_max} num_rows={int(view.active_slots.shape[0])}" + f"row_min={row_min} row_max={row_max} num_rows={int(meta.active_slots.shape[0])}" ) - visible_len = int(view.context_lens.max().item()) if view.context_lens.numel() > 0 else 0 - if visible_len > int(view.active_slots.shape[1]): + visible_len = int(meta.context_lens.max().item()) if meta.context_lens.numel() > 0 else 0 + if visible_len > int(meta.active_slots.shape[1]): raise RuntimeError( "decode visible length exceeds Req_to_tokens width: " - f"visible_len={visible_len} req_to_tokens_width={int(view.active_slots.shape[1])}" + f"visible_len={visible_len} req_to_tokens_width={int(meta.active_slots.shape[1])}" ) - visible_slots = view.active_slots.index_select(0, rows)[:, :visible_len] + visible_slots = meta.active_slots.index_select(0, rows)[:, :visible_len] pos = torch.arange(visible_len, device=visible_slots.device)[None, :] - valid_pos = pos < view.context_lens[:, None] - slot_cap = int(view.k_cache.shape[0]) + valid_pos = pos < meta.context_lens[:, None] + slot_cap = int(payload.k_cache.shape[0]) bad = ((visible_slots < 0) | (visible_slots >= slot_cap)) & valid_pos if bool(bad.any().item()): loc = bad.nonzero(as_tuple=False)[0] @@ -415,5 +456,5 @@ def _debug_check_decode_bounds(self, view: DecodeComputeView): raise RuntimeError( "decode physical slot out of bounds before attention: " f"batch={bad_b} req_row={bad_req_row} pos={bad_pos} " - f"slot={bad_slot} slot_cap={slot_cap} context_len={int(view.context_lens[bad_b].item())}" + f"slot={bad_slot} slot_cap={slot_cap} context_len={int(meta.context_lens[bad_b].item())}" ) diff --git a/tests/test_fake_attention_backend.py b/tests/test_fake_attention_backend.py index eb183f78..97ff5610 100644 --- a/tests/test_fake_attention_backend.py +++ b/tests/test_fake_attention_backend.py @@ -4,7 +4,13 @@ import torch -from sparsevllm.engine.cache_manager import DecodeComputeView, PrefillComputeView +from sparsevllm.engine.cache_manager import ( + AttentionViewMeta, + DecodeComputeView, + ExplicitKVPayload, + MlaLatentPayload, + PrefillComputeView, +) from sparsevllm.layers.attention_backend import TritonAttentionBackend @@ -31,24 +37,32 @@ def tearDown(self): def _make_prefill_view(self, *, attn_score=None): return PrefillComputeView( - k_cache=torch.ones(8, 2, 4), - v_cache=torch.ones(8, 2, 4), - active_slots=torch.tensor([[0, 1, 2]], dtype=torch.int32), - req_indices=torch.tensor([0], dtype=torch.int32), - context_lens=torch.tensor([3], dtype=torch.int32), - attn_score=attn_score, - max_context_len=3, + meta=AttentionViewMeta( + active_slots=torch.tensor([[0, 1, 2]], dtype=torch.int32), + req_indices=torch.tensor([0], dtype=torch.int32), + context_lens=torch.tensor([3], dtype=torch.int32), + attn_score=attn_score, + max_context_len=3, + ), + payload=ExplicitKVPayload( + k_cache=torch.ones(8, 2, 4), + v_cache=torch.ones(8, 2, 4), + ), ) def _make_decode_view(self, *, attn_score=None): return DecodeComputeView( - k_cache=torch.ones(8, 2, 4), - v_cache=torch.ones(8, 2, 4), - active_slots=torch.tensor([[0, 1, 2]], dtype=torch.int32), - req_indices=torch.tensor([0], dtype=torch.int32), - context_lens=torch.tensor([3], dtype=torch.int32), - attn_score=attn_score, - max_context_len=3, + meta=AttentionViewMeta( + active_slots=torch.tensor([[0, 1, 2]], dtype=torch.int32), + req_indices=torch.tensor([0], dtype=torch.int32), + context_lens=torch.tensor([3], dtype=torch.int32), + attn_score=attn_score, + max_context_len=3, + ), + payload=ExplicitKVPayload( + k_cache=torch.ones(8, 2, 4), + v_cache=torch.ones(8, 2, 4), + ), ) def test_fake_prefill_returns_zeros_and_skips_kernel(self): @@ -160,13 +174,17 @@ def test_debug_decode_bounds_checks_flash_attn_contiguous(self): os.environ["SVLLM_DEBUG_DECODE_BOUNDS"] = "1" q = torch.zeros(1, 2, 4) view = DecodeComputeView( - k_cache=torch.ones(3, 2, 4), - v_cache=torch.ones(3, 2, 4), - active_slots=torch.tensor([[0, 1, 99]], dtype=torch.int32), - req_indices=torch.tensor([0], dtype=torch.int32), - context_lens=torch.tensor([3], dtype=torch.int32), - max_context_len=3, - backend="flash_attn_contiguous", + meta=AttentionViewMeta( + active_slots=torch.tensor([[0, 1, 99]], dtype=torch.int32), + req_indices=torch.tensor([0], dtype=torch.int32), + context_lens=torch.tensor([3], dtype=torch.int32), + max_context_len=3, + ), + payload=ExplicitKVPayload( + k_cache=torch.ones(3, 2, 4), + v_cache=torch.ones(3, 2, 4), + backend="flash_attn_contiguous", + ), ) with self.assertRaisesRegex(RuntimeError, "decode physical slot out of bounds"): @@ -231,6 +249,47 @@ def stage2(mid_o, mid_o_logexpsum, context_lens, out, block_seq): self.assertTrue(torch.equal(out, torch.zeros_like(q))) self.assertTrue(torch.equal(attn_score, torch.full_like(attn_score, -1e20))) + def test_prefill_rejects_mla_payload_before_kernel(self): + explicit_view = self._make_prefill_view() + view = PrefillComputeView( + meta=explicit_view.meta, + payload=MlaLatentPayload( + latent_cache=torch.empty(8, 1, 512), + rope_cache=torch.empty(8, 1, 64), + ), + ) + + with self.assertRaisesRegex(TypeError, "requires ExplicitKVPayload"): + TritonAttentionBackend().run_prefill( + torch.empty(3, 2, 4), + view, + b_start_loc=torch.tensor([0], dtype=torch.int32), + chunk_lens=torch.tensor([3], dtype=torch.int32), + max_input_len=3, + ) + + def test_decode_rejects_mla_payload_before_kernel(self): + explicit_view = self._make_decode_view() + view = DecodeComputeView( + meta=explicit_view.meta, + payload=MlaLatentPayload( + latent_cache=torch.empty(8, 1, 512), + rope_cache=torch.empty(8, 1, 64), + ), + ) + + with self.assertRaisesRegex(TypeError, "requires ExplicitKVPayload"): + TritonAttentionBackend().run_decode( + torch.empty(1, 2, 4), + view, + mid_o=torch.empty(1, 2, 1, 4), + mid_o_logexpsum=torch.empty(1, 2, 1), + max_len_in_batch=3, + block_seq=256, + num_heads=2, + num_kv_heads=1, + ) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_h2o_cache_manager.py b/tests/test_h2o_cache_manager.py index bddf13e2..2a7530b9 100644 --- a/tests/test_h2o_cache_manager.py +++ b/tests/test_h2o_cache_manager.py @@ -9,7 +9,9 @@ import torch from sparsevllm.engine.cache_manager.base import ( + AttentionViewMeta, CacheManager, + ExplicitKVPayload, LayerBatchStates, PrefillComputeView, ) @@ -360,12 +362,16 @@ def test_h2o_prefill_score_collection_accumulates_in_physical_coordinates(): seq = _seq(0, 20, prefilled=8, chunk=2) manager._h2o_scores[(0, 0)] = torch.tensor([1.0, 2.0, 3.0, 4.0]) view = PrefillComputeView( - k_cache=torch.empty((16, 1, 1)), - v_cache=torch.empty((16, 1, 1)), - active_slots=manager.buffer_req_to_token_slots[0], - req_indices=torch.tensor([0], dtype=torch.int32), - context_lens=torch.tensor([6], dtype=torch.int32), - max_context_len=6, + meta=AttentionViewMeta( + active_slots=manager.buffer_req_to_token_slots[0], + req_indices=torch.tensor([0], dtype=torch.int32), + context_lens=torch.tensor([6], dtype=torch.int32), + max_context_len=6, + ), + payload=ExplicitKVPayload( + k_cache=torch.empty((16, 1, 1)), + v_cache=torch.empty((16, 1, 1)), + ), ) set_context(is_prefill=True, cache_manager=manager, seqs=[seq]) @@ -402,12 +408,16 @@ def test_h2o_prefill_score_collection_rejects_misaligned_physical_view(): seq = _seq(0, 20, prefilled=8, chunk=2) manager._h2o_scores[(0, 0)] = torch.ones(4) view = PrefillComputeView( - k_cache=torch.empty((16, 1, 1)), - v_cache=torch.empty((16, 1, 1)), - active_slots=manager.buffer_req_to_token_slots[0], - req_indices=torch.tensor([0], dtype=torch.int32), - context_lens=torch.tensor([7], dtype=torch.int32), - max_context_len=7, + meta=AttentionViewMeta( + active_slots=manager.buffer_req_to_token_slots[0], + req_indices=torch.tensor([0], dtype=torch.int32), + context_lens=torch.tensor([7], dtype=torch.int32), + max_context_len=7, + ), + payload=ExplicitKVPayload( + k_cache=torch.empty((16, 1, 1)), + v_cache=torch.empty((16, 1, 1)), + ), ) set_context(is_prefill=True, cache_manager=manager, seqs=[seq]) diff --git a/tests/test_prefill_schedule_policy.py b/tests/test_prefill_schedule_policy.py index 969359a6..025f024f 100644 --- a/tests/test_prefill_schedule_policy.py +++ b/tests/test_prefill_schedule_policy.py @@ -3745,7 +3745,11 @@ def compress(layer_idx): self.assertEqual(int(manager.row_seq_lens[row_idx]), 8) def test_deltakv_sparse_decode_backend_controls_fa2_view(self): - from sparsevllm.engine.cache_manager import DecodeComputeView + from sparsevllm.engine.cache_manager import ( + AttentionViewMeta, + DecodeComputeView, + ExplicitKVPayload, + ) from sparsevllm.engine.cache_manager.deltakv_base import DeltaKVCacheTritonManagerV4 from sparsevllm.engine.cache_manager.deltakv_less_memory import DeltaKVLessMemoryCacheManager @@ -3771,12 +3775,16 @@ def test_deltakv_sparse_decode_backend_controls_fa2_view(self): manager.deltakv_layer_to_idx = {1: 0} manager.has_prefill_staging_view = lambda layer_idx, active=staging_active: active view = DecodeComputeView( - k_cache=torch.empty((2, 1, 4), dtype=torch.float32), - v_cache=torch.empty((2, 1, 4), dtype=torch.float32), - active_slots=torch.tensor([[0, 1]], dtype=torch.int32), - req_indices=selection.req_indices, - context_lens=selection.context_lens, - backend="dense", + meta=AttentionViewMeta( + active_slots=torch.tensor([[0, 1]], dtype=torch.int32), + req_indices=selection.req_indices, + context_lens=selection.context_lens, + ), + payload=ExplicitKVPayload( + k_cache=torch.empty((2, 1, 4), dtype=torch.float32), + v_cache=torch.empty((2, 1, 4), dtype=torch.float32), + backend="dense", + ), ) with patch.object(DeltaKVCacheTritonManagerV4, "build_decode_compute_view", return_value=view): @@ -3789,7 +3797,8 @@ def test_deltakv_sparse_decode_backend_controls_fa2_view(self): num_kv_heads=1, ) - self.assertEqual(out.backend, expected) + self.assertIsInstance(out.payload, ExplicitKVPayload) + self.assertEqual(out.payload.backend, expected) def test_static_decode_resets_deltakv_view_cache_before_validation(self): manager = object.__new__(DeltaKVCacheManager) diff --git a/tests/test_qwen35_hd256_decode_routing.py b/tests/test_qwen35_hd256_decode_routing.py index b9ba58ba..09daeb27 100644 --- a/tests/test_qwen35_hd256_decode_routing.py +++ b/tests/test_qwen35_hd256_decode_routing.py @@ -3,7 +3,11 @@ import torch -from sparsevllm.engine.cache_manager import DecodeComputeView +from sparsevllm.engine.cache_manager import ( + AttentionViewMeta, + DecodeComputeView, + ExplicitKVPayload, +) from sparsevllm.layers.attention_backend import TritonAttentionBackend from sparsevllm.triton_kernel.flash_decoding_stage2 import flash_decode_stage2 from sparsevllm.triton_kernel.gqa_flash_decoding_stage1 import ( @@ -20,13 +24,14 @@ def _make_view(self, *, head_dim: int, attn_score=None): k_cache = torch.zeros(4, 4, head_dim, dtype=torch.float32) v_cache = torch.zeros_like(k_cache) return DecodeComputeView( - k_cache=k_cache, - v_cache=v_cache, - active_slots=active_slots, - req_indices=req_indices, - context_lens=context_lens, - attn_score=attn_score, - max_context_len=3, + meta=AttentionViewMeta( + active_slots=active_slots, + req_indices=req_indices, + context_lens=context_lens, + attn_score=attn_score, + max_context_len=3, + ), + payload=ExplicitKVPayload(k_cache=k_cache, v_cache=v_cache), ) def test_head_dim_256_uses_grouped_gqa_stage1_and_unified_stage2(self): diff --git a/tests/test_rkv_skipkv_methods.py b/tests/test_rkv_skipkv_methods.py index ecaeac7a..1f6d6789 100644 --- a/tests/test_rkv_skipkv_methods.py +++ b/tests/test_rkv_skipkv_methods.py @@ -7,7 +7,12 @@ import torch from sparsevllm.config import Config, RuntimeLayout -from sparsevllm.engine.cache_manager.base import LayerBatchStates +from sparsevllm.engine.cache_manager.base import ( + AttentionViewMeta, + ExplicitKVPayload, + LayerBatchStates, + PrefillComputeView, +) from sparsevllm.engine.cache_manager.rkv import RKVCacheManager from sparsevllm.engine.cache_manager.skipkv import ( SkipKVCacheManager, @@ -502,9 +507,16 @@ def test_rkv_query_cache_tracks_observation_tokens_not_interval(self): manager._rkv_query_positions = [torch.full((1, 3), -1, dtype=torch.int32)] q_prefill = torch.arange(10, dtype=torch.float32).view(5, 1, 2) - view = SimpleNamespace( - req_indices=torch.tensor([0], dtype=torch.int32), - context_lens=torch.tensor([5], dtype=torch.int32), + view = PrefillComputeView( + meta=AttentionViewMeta( + active_slots=torch.empty((1, 0), dtype=torch.int32), + req_indices=torch.tensor([0], dtype=torch.int32), + context_lens=torch.tensor([5], dtype=torch.int32), + ), + payload=ExplicitKVPayload( + k_cache=torch.empty((0, 1, 2)), + v_cache=torch.empty((0, 1, 2)), + ), ) manager.record_prefill_query( 0, From 7eea70cc05abce02dac9c213986f027ec81b40e6 Mon Sep 17 00:00:00 2001 From: Codex Date: Tue, 4 Aug 2026 22:24:49 +0800 Subject: [PATCH 03/48] feat: add mla cache storage --- src/sparsevllm/configs/model.py | 39 ++ src/sparsevllm/configs/runtime.py | 1 + .../engine/cache_manager/__init__.py | 6 + src/sparsevllm/engine/cache_manager/base.py | 140 +++++- .../engine/cache_manager/standard.py | 130 +++++- .../engine/cache_manager/storage/__init__.py | 51 +++ .../engine/cache_manager/storage/base.py | 43 ++ .../cache_manager/storage/explicit_kv.py | 153 +++++++ .../cache_manager/storage/mla_latent.py | 160 +++++++ tests/test_attention_cache_storage.py | 424 ++++++++++++++++++ 10 files changed, 1113 insertions(+), 34 deletions(-) create mode 100644 src/sparsevllm/engine/cache_manager/storage/__init__.py create mode 100644 src/sparsevllm/engine/cache_manager/storage/base.py create mode 100644 src/sparsevllm/engine/cache_manager/storage/explicit_kv.py create mode 100644 src/sparsevllm/engine/cache_manager/storage/mla_latent.py create mode 100644 tests/test_attention_cache_storage.py diff --git a/src/sparsevllm/configs/model.py b/src/sparsevllm/configs/model.py index cf4a598d..635ac201 100644 --- a/src/sparsevllm/configs/model.py +++ b/src/sparsevllm/configs/model.py @@ -5,9 +5,11 @@ from types import SimpleNamespace from typing import Any +import torch from transformers import AutoConfig from sparsevllm.method_registry import ( + normalize_sparse_method, validate_model_runtime_compatibility, validate_sparse_method_assets, ) @@ -62,6 +64,42 @@ def _extract_text_config(config: Any) -> Any: return text_config +def _resolve_attention_cache_layout(config, model_type: str) -> None: + config.attention_cache_layout = ( + "mla_latent" if model_type == "glm4_moe_lite" else "explicit_kv" + ) + if config.attention_cache_layout == "explicit_kv": + return + if normalize_sparse_method(config.vllm_sparse_method): + raise NotImplementedError( + "GLM-4.7-Flash latent MLA supports only vanilla attention." + ) + if config.enable_prefix_caching or config.enable_prefix_cache_offload: + raise NotImplementedError( + "GLM-4.7-Flash latent MLA does not support prefix caching or offload." + ) + if config.decode_cuda_graph or not config.enforce_eager: + raise NotImplementedError( + "GLM-4.7-Flash latent MLA requires eager execution." + ) + if config.quantization_config.enabled: + raise NotImplementedError( + "GLM-4.7-Flash latent MLA does not support quantized checkpoints." + ) + if config.hf_config.torch_dtype != torch.bfloat16: + raise NotImplementedError( + "GLM-4.7-Flash latent MLA requires BF16 checkpoints." + ) + if ( + int(config_get(config.hf_config, "kv_lora_rank", 0)) != 512 + or int(config_get(config.hf_config, "qk_rope_head_dim", 0)) != 64 + ): + raise NotImplementedError( + "GLM-4.7-Flash latent MLA requires kv_lora_rank=512 and " + "qk_rope_head_dim=64." + ) + + def _finalize_model_config(config, model_spec: ModelSpec) -> None: config.runtime_layout = RuntimeLayout.from_config( config.hf_config, @@ -132,6 +170,7 @@ def load_and_validate_model(config) -> None: "Tiny random mode does not support quantized model weights." ) setattr(config.hf_config, "quantization_config", config.quantization_config) + _resolve_attention_cache_layout(config, model_type) validate_checkpoint( model_type, outer_config=config.outer_hf_config, diff --git a/src/sparsevllm/configs/runtime.py b/src/sparsevllm/configs/runtime.py index 94b3895a..0c1a1467 100644 --- a/src/sparsevllm/configs/runtime.py +++ b/src/sparsevllm/configs/runtime.py @@ -74,6 +74,7 @@ class Config( hf_config: AutoConfig | None = None outer_hf_config: Any | None = None runtime_layout: RuntimeLayout | None = None + attention_cache_layout: str = field(default="explicit_kv", init=False) quantization_config: QuantizationConfig = field(default_factory=QuantizationConfig.disabled) model_spec: ModelSpec = field(init=False, repr=False) parallel_topology: ParallelTopology = field(init=False, repr=False) diff --git a/src/sparsevllm/engine/cache_manager/__init__.py b/src/sparsevllm/engine/cache_manager/__init__.py index 132dd24f..c27f80fb 100644 --- a/src/sparsevllm/engine/cache_manager/__init__.py +++ b/src/sparsevllm/engine/cache_manager/__init__.py @@ -1,25 +1,31 @@ from __future__ import annotations from .base import ( + AttentionCacheWrite, AttentionPayload, AttentionViewMeta, CacheManager, DecodeComputeView, ExplicitKVPayload, + ExplicitKVWrite, LayerBatchStates, MlaLatentPayload, + MlaLatentWrite, PrefillComputeView, SparseSelection, ) __all__ = [ + "AttentionCacheWrite", "AttentionPayload", "AttentionViewMeta", "CacheManager", "DecodeComputeView", "ExplicitKVPayload", + "ExplicitKVWrite", "LayerBatchStates", "MlaLatentPayload", + "MlaLatentWrite", "PrefillComputeView", "SparseSelection", "StandardCacheManager", diff --git a/src/sparsevllm/engine/cache_manager/base.py b/src/sparsevllm/engine/cache_manager/base.py index 5fab60f6..a698424e 100644 --- a/src/sparsevllm/engine/cache_manager/base.py +++ b/src/sparsevllm/engine/cache_manager/base.py @@ -174,6 +174,25 @@ class MlaLatentPayload: AttentionPayload = ExplicitKVPayload | MlaLatentPayload +@dataclass(frozen=True) +class ExplicitKVWrite: + """Current-token key/value tensors to persist.""" + + key: torch.Tensor + value: torch.Tensor + + +@dataclass(frozen=True) +class MlaLatentWrite: + """Current-token latent and RoPE tensors to persist.""" + + latent: torch.Tensor + rope: torch.Tensor + + +AttentionCacheWrite = ExplicitKVWrite | MlaLatentWrite + + @dataclass(frozen=True) class DecodeComputeView: """Decode metadata paired with exactly one physical payload layout.""" @@ -436,7 +455,7 @@ def _get_available_slots_info(self) -> tuple[int, int]: + current - recurrent_explicit_deduction ) - slot_bytes_per_layer = 2 * self.num_kv_heads * self.head_dim * dtype_size + slot_bytes_per_layer = self.attention_cache_bytes_per_slot_per_layer() recurrent_bytes_per_block = int( getattr(config, "prefix_recurrent_bytes_per_block", 0) or 0 @@ -527,6 +546,11 @@ def _get_available_slots_info(self) -> tuple[int, int]: return available_memory, slot_bytes_per_layer + def attention_cache_bytes_per_slot_per_layer(self) -> int: + """Persistent attention-cache bytes for one token in one KV layer.""" + dtype_size = self._cache_slot_dtype_size() + return int(2 * self.num_kv_heads * self.head_dim * dtype_size) + def _kv_allocation_bytes_per_prefix_block( self, slot_bytes_per_layer: int, @@ -595,6 +619,23 @@ def _store_layer_kv( store_kvcache(k, v, k_cache, v_cache, slot_mapping) return slot_mapping + def store_attention_payload( + self, + layer_idx: int, + payload: AttentionCacheWrite, + ) -> torch.Tensor: + """Store one layer's current-token payload using the configured layout.""" + if not isinstance(payload, ExplicitKVWrite): + raise TypeError( + f"{type(self).__name__} supports only ExplicitKVWrite stores, got " + f"{type(payload).__name__}." + ) + return self._store_layer_kv( + layer_idx, + payload.key, + payload.value, + ) + def save_raw_kv_if_needed( self, layer_idx: int, @@ -617,7 +658,10 @@ def save_rope_kv_if_needed( k_post_rope=k_post_rope, v=v, ) - slot_mapping = self._store_layer_kv(layer_idx, store_k, store_v) + slot_mapping = self.store_attention_payload( + layer_idx, + ExplicitKVWrite(key=store_k, value=store_v), + ) self.on_kv_stored( layer_idx, store_k, @@ -639,6 +683,31 @@ def get_layer_compute_view( k_cache, v_cache = self.get_layer_kv_cache(layer_idx) return k_cache, v_cache, active_slots, req_indices, context_lens + def get_layer_compute_payload( + self, + layer_idx: int, + active_slots: torch.Tensor, + req_indices: torch.Tensor, + context_lens: torch.Tensor, + selection: SparseSelection | None = None, + ) -> tuple[AttentionPayload, torch.Tensor, torch.Tensor, torch.Tensor]: + """Return the physical payload and logical coordinates for decode.""" + k_cache, v_cache, active_slots, req_indices, context_lens = ( + self.get_layer_compute_view( + layer_idx, + active_slots, + req_indices, + context_lens, + selection, + ) + ) + return ( + ExplicitKVPayload(k_cache=k_cache, v_cache=v_cache), + active_slots, + req_indices, + context_lens, + ) + def get_prefill_compute_view( self, layer_idx: int, @@ -659,6 +728,35 @@ def get_prefill_compute_view( selection, ) + def get_prefill_compute_payload( + self, + layer_idx: int, + k_current: torch.Tensor, + v_current: torch.Tensor, + selection: SparseSelection, + active_slots: torch.Tensor, + req_indices: torch.Tensor, + context_lens: torch.Tensor, + ) -> tuple[AttentionPayload, torch.Tensor, torch.Tensor, torch.Tensor]: + """Return the physical payload and logical coordinates for prefill.""" + k_cache, v_cache, active_slots, req_indices, context_lens = ( + self.get_prefill_compute_view( + layer_idx, + k_current, + v_current, + selection, + active_slots, + req_indices, + context_lens, + ) + ) + return ( + ExplicitKVPayload(k_cache=k_cache, v_cache=v_cache), + active_slots, + req_indices, + context_lens, + ) + def _default_active_slots_for_selection(self, layer_idx: int, selection: SparseSelection) -> torch.Tensor: if selection.active_slots is not None: return selection.active_slots @@ -684,7 +782,7 @@ def build_prefill_compute_view( active_slots = self._default_active_slots_for_selection(layer_idx, selection) req_indices = selection.req_indices context_lens = selection.context_lens - k_cache, v_cache, active_slots, req_indices, context_lens = self.get_prefill_compute_view( + payload, active_slots, req_indices, context_lens = self.get_prefill_compute_payload( layer_idx, k_current, v_current, @@ -702,7 +800,7 @@ def build_prefill_compute_view( max_context_len=selection.max_context_len, temp_slots=temp_slots, ), - payload=ExplicitKVPayload(k_cache=k_cache, v_cache=v_cache), + payload=payload, ) def collect_prefill_attention_score( @@ -933,12 +1031,14 @@ def build_decode_compute_view( num_heads=num_heads, num_kv_heads=num_kv_heads, ) - k_cache, v_cache, active_slots, req_indices, context_lens = self.get_layer_compute_view( - layer_idx, - active_slots, - req_indices, - context_lens, - selection, + payload, active_slots, req_indices, context_lens = ( + self.get_layer_compute_payload( + layer_idx, + active_slots, + req_indices, + context_lens, + selection, + ) ) return DecodeComputeView( meta=AttentionViewMeta( @@ -948,7 +1048,7 @@ def build_decode_compute_view( attn_score=selection.attn_score, max_context_len=selection.max_context_len, ), - payload=ExplicitKVPayload(k_cache=k_cache, v_cache=v_cache), + payload=payload, ) def get_decode_block_seq(self, layer_idx: int, default: int) -> int: @@ -1285,7 +1385,8 @@ def debug_state_summary(self) -> dict[str, Any]: } def _cache_slot_dtype_size(self) -> int: - dtype = getattr(self.hf_config, "torch_dtype", torch.float16) + hf_config = getattr(self, "hf_config", getattr(self.config, "hf_config", None)) + dtype = getattr(hf_config, "torch_dtype", torch.float16) if not isinstance(dtype, torch.dtype): dtype = torch.float16 return int(torch.tensor([], dtype=dtype).element_size()) @@ -1350,8 +1451,14 @@ def visit(path: str, value): for field_name, item in value.__dict__.items(): yield from visit(f"{path}.{field_name}", item) + storage = getattr(self, "attention_cache_storage", None) + if storage is not None: + layout = getattr(getattr(storage, "layout", None), "value", "unknown") + for index, tensor in enumerate(storage.accounting_tensors()): + yield f"attention_cache_storage.{layout}.{index}_cache", tensor + for name, value in self.__dict__.items(): - if name in {"config", "hf_config"}: + if name in {"config", "hf_config", "attention_cache_storage"}: continue yield from visit(name, value) @@ -1376,8 +1483,11 @@ def _logical_live_kv_bytes(self) -> int: live_tokens = int(row_seq_lens.sum()) except Exception: return 0 - dtype_size = self._cache_slot_dtype_size() - return int(live_tokens * self.num_kv_layers * 2 * self.num_kv_heads * self.head_dim * dtype_size) + return int( + live_tokens + * self.num_kv_layers + * self.attention_cache_bytes_per_slot_per_layer() + ) def memory_accounting(self) -> dict[str, Any]: """Return read-only tensor memory accounting for regression gates. diff --git a/src/sparsevllm/engine/cache_manager/standard.py b/src/sparsevllm/engine/cache_manager/standard.py index a4693903..d603e34a 100644 --- a/src/sparsevllm/engine/cache_manager/standard.py +++ b/src/sparsevllm/engine/cache_manager/standard.py @@ -23,13 +23,23 @@ from sparsevllm.utils.profiler import profiler from sparsevllm.platforms import device_runtime -from .base import CacheManager, LayerBatchStates, SparseSelection +from .base import ( + AttentionCacheWrite, + AttentionPayload, + CacheManager, + LayerBatchStates, + SparseSelection, +) from .prefix_cache_mixin import PrefixCacheMixin from .prefix_offload import ( PinnedPrefixKVPool, PrefixH2DOperation, StandardPrefixOffloadController, ) +from .storage import ( + ExplicitKVStorage, + create_attention_cache_storage, +) @dataclass @@ -68,6 +78,11 @@ class StandardCacheManager(PrefixCacheMixin, CacheManager): def __init__(self, config: Config, parallel_context: ParallelContext): super().__init__(config, parallel_context) + self.attention_cache_storage = create_attention_cache_storage( + config, + num_kv_heads=self.num_kv_heads, + head_dim=self.head_dim, + ) self.allocate_kv_cache() num_slots = config.num_kvcache_slots @@ -124,13 +139,12 @@ def _init_prefix_offload(self) -> None: host_size_gb = getattr(self.config, "prefix_cache_host_size_gb", None) if host_size_gb is None: raise RuntimeError("Prefix cache offload requires prefix_cache_host_size_gb.") + storage = self._require_explicit_storage("Prefix cache offload") + kv_cache = storage.cache bytes_per_block = int( self.prefix_cache_block_size * self.num_kv_layers - * 2 - * self.num_kv_heads - * self.head_dim - * self.kv_cache.element_size() + * storage.bytes_per_slot_per_layer() ) host_bytes = int(float(host_size_gb) * (1024**3)) host_capacity_blocks = host_bytes // bytes_per_block @@ -148,13 +162,13 @@ def _init_prefix_offload(self) -> None: capacity_blocks=host_capacity_blocks, num_layers=self.num_kv_layers, block_size=self.prefix_cache_block_size, - num_kv_heads=self.num_kv_heads, - head_dim=self.head_dim, - dtype=self.kv_cache.dtype, + num_kv_heads=storage.num_kv_heads, + head_dim=storage.head_dim, + dtype=storage.dtype, ) self.prefix_offload_controller = StandardPrefixOffloadController( prefix_cache=self.prefix_cache, - kv_cache=self.kv_cache, + kv_cache=kv_cache, host_pool=host_pool, block_size=self.prefix_cache_block_size, device=self.device, @@ -179,26 +193,101 @@ def allocate_kv_cache(self): logger.info( f"Standard Mode: Each layer can accommodate {self.config.num_kvcache_slots} tokens." ) - self.kv_cache = torch.empty( - 2, - num_layers, - self.config.num_kvcache_slots, - self.num_kv_heads, - self.head_dim, - dtype=self.hf_config.torch_dtype, + self.attention_cache_storage.allocate( + num_layers=num_layers, + num_slots=self.config.num_kvcache_slots, device=self.device, ) + self.kv_cache = ( + self.attention_cache_storage.kv_cache + if isinstance(self.attention_cache_storage, ExplicitKVStorage) + else None + ) + + def attention_cache_bytes_per_slot_per_layer(self) -> int: + storage = getattr(self, "attention_cache_storage", None) + if storage is None: + return super().attention_cache_bytes_per_slot_per_layer() + return int(storage.bytes_per_slot_per_layer()) + + def _require_explicit_storage(self, operation: str) -> ExplicitKVStorage: + storage = self.attention_cache_storage + if not isinstance(storage, ExplicitKVStorage): + raise TypeError( + f"{operation} requires ExplicitKVStorage, got " + f"{type(storage).__name__}." + ) + return storage def get_layer_batch_states(self, layer_idx: int) -> LayerBatchStates: return self.layer_batch_state def get_layer_kv_cache(self, layer_idx: int) -> tuple[torch.Tensor, torch.Tensor]: kv_idx = self.kv_layer_index(layer_idx) - return self.kv_cache[0, kv_idx], self.kv_cache[1, kv_idx] + payload = self._require_explicit_storage("get_layer_kv_cache").layer_payload(kv_idx) + return payload.k_cache, payload.v_cache def get_layer_store_view(self, layer_idx: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - kv_idx = self.kv_layer_index(layer_idx) - return self.kv_cache[0, kv_idx], self.kv_cache[1, kv_idx], self.layer_batch_state.slot_mapping + k_cache, v_cache = self.get_layer_kv_cache(layer_idx) + return k_cache, v_cache, self.layer_batch_state.slot_mapping + + def store_attention_payload( + self, + layer_idx: int, + payload: AttentionCacheWrite, + ) -> torch.Tensor: + slot_mapping = self.layer_batch_state.slot_mapping + if slot_mapping is None: + raise RuntimeError( + f"Attention cache store requires slot_mapping at layer={layer_idx}." + ) + self.attention_cache_storage.store( + self.kv_layer_index(layer_idx), + slot_mapping, + payload, + ) + return slot_mapping + + def _validate_attention_slot_mapping(self, slot_mapping: torch.Tensor) -> None: + storage = getattr(self, "attention_cache_storage", None) + if storage is not None: + storage.validate_slot_mapping(slot_mapping) + + def get_layer_compute_payload( + self, + layer_idx: int, + active_slots: torch.Tensor, + req_indices: torch.Tensor, + context_lens: torch.Tensor, + selection: SparseSelection | None = None, + ) -> tuple[AttentionPayload, torch.Tensor, torch.Tensor, torch.Tensor]: + del selection + return ( + self.attention_cache_storage.layer_payload( + self.kv_layer_index(layer_idx) + ), + active_slots, + req_indices, + context_lens, + ) + + def get_prefill_compute_payload( + self, + layer_idx: int, + k_current: torch.Tensor, + v_current: torch.Tensor, + selection: SparseSelection, + active_slots: torch.Tensor, + req_indices: torch.Tensor, + context_lens: torch.Tensor, + ) -> tuple[AttentionPayload, torch.Tensor, torch.Tensor, torch.Tensor]: + del k_current, v_current, selection + return self.get_layer_compute_payload( + layer_idx, + active_slots, + req_indices, + context_lens, + ) def get_layer_compute_tensors(self, layer_idx: int, selection: SparseSelection | None = None): del selection @@ -1283,6 +1372,7 @@ def _prepare_prefill(self, seqs: list[Sequence]): self.layer_batch_state.context_lens = context_lens self.layer_batch_state.max_context_len = max(context_lens_list) if context_lens_list else 0 self.layer_batch_state.req_indices = req_indices_tensor + self._validate_attention_slot_mapping(slot_mapping) if log_level == 'DEBUG': logger.debug(f'{context_lens_list=} {req_indices=} {slot_mapping[:10].tolist()=} {slot_mapping[-10:].tolist()=}') @@ -1319,6 +1409,7 @@ def _prepare_decode(self, seqs: list[Sequence]): self.layer_batch_state.context_lens = context_lens self.layer_batch_state.max_context_len = int(max(self.row_seq_lens[row_indices])) if row_indices else 0 self.layer_batch_state.req_indices = req_indices + self._validate_attention_slot_mapping(slot_mapping) if log_level == 'DEBUG': logger.debug(f'{slot_mapping=} {context_lens.tolist()=} {slot_mapping[:10]=} {slot_mapping[-10:]=}') @@ -1414,5 +1505,6 @@ def prepare_decode_static( self.layer_batch_state.context_lens = context_lens self.layer_batch_state.max_context_len = int(real_context_lens.max()) if real_batch_size > 0 else 0 self.layer_batch_state.req_indices = req_indices + self._validate_attention_slot_mapping(slot_mapping) return input_ids, positions, None diff --git a/src/sparsevllm/engine/cache_manager/storage/__init__.py b/src/sparsevllm/engine/cache_manager/storage/__init__.py new file mode 100644 index 00000000..8c5f20d6 --- /dev/null +++ b/src/sparsevllm/engine/cache_manager/storage/__init__.py @@ -0,0 +1,51 @@ +from __future__ import annotations + +from typing import Any + +from .base import AttentionCacheStorage, CacheLayout +from .explicit_kv import ExplicitKVStorage +from .mla_latent import MlaLatentStorage + + +def create_attention_cache_storage( + config: Any, + *, + num_kv_heads: int, + head_dim: int, +) -> AttentionCacheStorage: + configured_layout = getattr(config, "attention_cache_layout", None) + if configured_layout is None: + model_type = str(getattr(config.hf_config, "model_type", "") or "") + configured_layout = ( + CacheLayout.MLA_LATENT.value + if model_type == "glm4_moe_lite" + else CacheLayout.EXPLICIT_KV.value + ) + layout = ( + configured_layout + if isinstance(configured_layout, CacheLayout) + else CacheLayout(str(configured_layout)) + ) + dtype = config.hf_config.torch_dtype + if layout is CacheLayout.EXPLICIT_KV: + return ExplicitKVStorage( + num_kv_heads=num_kv_heads, + head_dim=head_dim, + dtype=dtype, + ) + if layout is CacheLayout.MLA_LATENT: + return MlaLatentStorage( + kv_lora_rank=int(config.hf_config.kv_lora_rank), + rope_dim=int(config.hf_config.qk_rope_head_dim), + dtype=dtype, + ) + raise AssertionError(f"Unhandled attention cache layout: {layout!r}") + + +__all__ = [ + "AttentionCacheStorage", + "CacheLayout", + "ExplicitKVStorage", + "MlaLatentStorage", + "create_attention_cache_storage", +] diff --git a/src/sparsevllm/engine/cache_manager/storage/base.py b/src/sparsevllm/engine/cache_manager/storage/base.py new file mode 100644 index 00000000..f64726ef --- /dev/null +++ b/src/sparsevllm/engine/cache_manager/storage/base.py @@ -0,0 +1,43 @@ +from __future__ import annotations + +from enum import Enum +from typing import Protocol, runtime_checkable + +import torch + +from ..base import AttentionCacheWrite, AttentionPayload + + +class CacheLayout(str, Enum): + EXPLICIT_KV = "explicit_kv" + MLA_LATENT = "mla_latent" + + +@runtime_checkable +class AttentionCacheStorage(Protocol): + """Physical attention-cache storage owned by a cache manager.""" + + layout: CacheLayout + + def allocate( + self, + *, + num_layers: int, + num_slots: int, + device: torch.device, + ) -> None: ... + + def layer_payload(self, layer_idx: int) -> AttentionPayload: ... + + def validate_slot_mapping(self, slot_mapping: torch.Tensor) -> None: ... + + def store( + self, + layer_idx: int, + slot_mapping: torch.Tensor, + payload: AttentionCacheWrite, + ) -> None: ... + + def bytes_per_slot_per_layer(self) -> int: ... + + def accounting_tensors(self) -> tuple[torch.Tensor, ...]: ... diff --git a/src/sparsevllm/engine/cache_manager/storage/explicit_kv.py b/src/sparsevllm/engine/cache_manager/storage/explicit_kv.py new file mode 100644 index 00000000..f2db9734 --- /dev/null +++ b/src/sparsevllm/engine/cache_manager/storage/explicit_kv.py @@ -0,0 +1,153 @@ +from __future__ import annotations + +import torch + +from sparsevllm.triton_kernel.store_kvcache import store_kvcache + +from ..base import AttentionCacheWrite, ExplicitKVPayload, ExplicitKVWrite +from .base import CacheLayout + + +class ExplicitKVStorage: + """The ordinary two-tensor K/V cache layout.""" + + layout = CacheLayout.EXPLICIT_KV + + def __init__( + self, + *, + num_kv_heads: int, + head_dim: int, + dtype: torch.dtype, + ) -> None: + self.num_kv_heads = int(num_kv_heads) + self.head_dim = int(head_dim) + self.dtype = dtype + if self.num_kv_heads <= 0 or self.head_dim <= 0: + raise ValueError( + "Explicit KV dimensions must be positive, got " + f"num_kv_heads={self.num_kv_heads} head_dim={self.head_dim}." + ) + self.kv_cache: torch.Tensor | None = None + + def allocate( + self, + *, + num_layers: int, + num_slots: int, + device: torch.device, + ) -> None: + num_layers = int(num_layers) + num_slots = int(num_slots) + if num_layers <= 0 or num_slots <= 0: + raise ValueError( + "Explicit KV allocation requires positive layers and slots, got " + f"num_layers={num_layers} num_slots={num_slots}." + ) + self.kv_cache = torch.empty( + 2, + num_layers, + num_slots, + self.num_kv_heads, + self.head_dim, + dtype=self.dtype, + device=device, + ) + + def _require_cache(self) -> torch.Tensor: + if self.kv_cache is None: + raise RuntimeError("Explicit KV storage has not been allocated.") + return self.kv_cache + + @property + def cache(self) -> torch.Tensor: + return self._require_cache() + + def layer_payload(self, layer_idx: int) -> ExplicitKVPayload: + cache = self._require_cache() + layer_idx = int(layer_idx) + if not 0 <= layer_idx < int(cache.shape[1]): + raise IndexError( + f"Explicit KV layer index {layer_idx} is outside [0, {int(cache.shape[1])})." + ) + return ExplicitKVPayload( + k_cache=cache[0, layer_idx], + v_cache=cache[1, layer_idx], + ) + + def validate_slot_mapping(self, slot_mapping: torch.Tensor) -> None: + cache = self._require_cache() + if slot_mapping.ndim != 1: + raise ValueError( + f"Explicit KV slot_mapping must be 1D, got {tuple(slot_mapping.shape)}." + ) + if slot_mapping.dtype != torch.int32: + raise TypeError( + f"Explicit KV slot_mapping must use torch.int32, got {slot_mapping.dtype}." + ) + if slot_mapping.device != cache.device: + raise ValueError( + "Explicit KV slot_mapping must share the cache device, got " + f"slots={slot_mapping.device} cache={cache.device}." + ) + + def store( + self, + layer_idx: int, + slot_mapping: torch.Tensor, + payload: AttentionCacheWrite, + ) -> None: + if not isinstance(payload, ExplicitKVWrite): + raise TypeError( + "ExplicitKVStorage.store requires ExplicitKVWrite, got " + f"{type(payload).__name__}." + ) + destination = self.layer_payload(layer_idx) + if payload.key.shape != payload.value.shape: + raise ValueError( + "Explicit KV store tensors must have equal shapes, got " + f"k={tuple(payload.key.shape)} v={tuple(payload.value.shape)}." + ) + expected_tail = (self.num_kv_heads, self.head_dim) + if payload.key.ndim != 3 or tuple(payload.key.shape[1:]) != expected_tail: + raise ValueError( + "Explicit KV store tensors must have shape [tokens, " + f"{self.num_kv_heads}, {self.head_dim}], got " + f"{tuple(payload.key.shape)}." + ) + if slot_mapping.shape != (int(payload.key.shape[0]),): + raise ValueError( + "Explicit KV slot_mapping must match the token dimension, got " + f"slots={tuple(slot_mapping.shape)} tokens={int(payload.key.shape[0])}." + ) + if payload.key.dtype != self.dtype or payload.value.dtype != self.dtype: + raise TypeError( + f"Explicit KV store requires dtype={self.dtype}, got " + f"k={payload.key.dtype} v={payload.value.dtype}." + ) + destination_device = destination.k_cache.device + if ( + payload.key.device != destination_device + or payload.value.device != destination_device + or slot_mapping.device != destination_device + ): + raise ValueError( + "Explicit KV store tensors must share the destination device, got " + f"destination={destination_device} key={payload.key.device} " + f"value={payload.value.device} slots={slot_mapping.device}." + ) + self.validate_slot_mapping(slot_mapping) + store_kvcache( + payload.key, + payload.value, + destination.k_cache, + destination.v_cache, + slot_mapping, + ) + + def bytes_per_slot_per_layer(self) -> int: + element_size = torch.tensor([], dtype=self.dtype).element_size() + return int(2 * self.num_kv_heads * self.head_dim * element_size) + + def accounting_tensors(self) -> tuple[torch.Tensor, ...]: + return (self._require_cache(),) diff --git a/src/sparsevllm/engine/cache_manager/storage/mla_latent.py b/src/sparsevllm/engine/cache_manager/storage/mla_latent.py new file mode 100644 index 00000000..3de9575d --- /dev/null +++ b/src/sparsevllm/engine/cache_manager/storage/mla_latent.py @@ -0,0 +1,160 @@ +from __future__ import annotations + +import torch + +from sparsevllm.triton_kernel.mla.copy_latent import ( + copy_latent_to_cache, + validate_copy_slot_mapping, +) +from sparsevllm.triton_kernel.mla.decode_stage1 import MLA_LATENT_DIM, MLA_ROPE_DIM + +from ..base import AttentionCacheWrite, MlaLatentPayload, MlaLatentWrite +from .base import CacheLayout + + +class MlaLatentStorage: + """Persistent latent and RoPE caches for GLM-style MLA.""" + + layout = CacheLayout.MLA_LATENT + + def __init__( + self, + *, + kv_lora_rank: int, + rope_dim: int, + dtype: torch.dtype, + ) -> None: + self.kv_lora_rank = int(kv_lora_rank) + self.rope_dim = int(rope_dim) + self.dtype = dtype + if self.kv_lora_rank != MLA_LATENT_DIM or self.rope_dim != MLA_ROPE_DIM: + raise ValueError( + "The vendored MLA storage kernel requires " + f"kv_lora_rank={MLA_LATENT_DIM} and rope_dim={MLA_ROPE_DIM}, got " + f"kv_lora_rank={self.kv_lora_rank} rope_dim={self.rope_dim}." + ) + if self.dtype != torch.bfloat16: + raise TypeError( + f"MLA latent storage v1 requires torch.bfloat16, got {self.dtype}." + ) + self.latent_cache: torch.Tensor | None = None + self.rope_cache: torch.Tensor | None = None + self._validated_slot_mapping_key: tuple[str, int, int] | None = None + self._validated_store_calls_remaining = 0 + + def allocate( + self, + *, + num_layers: int, + num_slots: int, + device: torch.device, + ) -> None: + num_layers = int(num_layers) + num_slots = int(num_slots) + if num_layers <= 0 or num_slots <= 0: + raise ValueError( + "MLA allocation requires positive layers and slots, got " + f"num_layers={num_layers} num_slots={num_slots}." + ) + self.latent_cache = torch.empty( + num_layers, + num_slots, + 1, + self.kv_lora_rank, + dtype=self.dtype, + device=device, + ) + self.rope_cache = torch.empty( + num_layers, + num_slots, + 1, + self.rope_dim, + dtype=self.dtype, + device=device, + ) + self._validated_slot_mapping_key = None + self._validated_store_calls_remaining = 0 + + def _require_caches(self) -> tuple[torch.Tensor, torch.Tensor]: + if self.latent_cache is None or self.rope_cache is None: + raise RuntimeError("MLA latent storage has not been allocated.") + return self.latent_cache, self.rope_cache + + def layer_payload(self, layer_idx: int) -> MlaLatentPayload: + latent_cache, rope_cache = self._require_caches() + layer_idx = int(layer_idx) + if not 0 <= layer_idx < int(latent_cache.shape[0]): + raise IndexError( + f"MLA layer index {layer_idx} is outside [0, {int(latent_cache.shape[0])})." + ) + return MlaLatentPayload( + latent_cache=latent_cache[layer_idx], + rope_cache=rope_cache[layer_idx], + ) + + @staticmethod + def _slot_mapping_key(slot_mapping: torch.Tensor) -> tuple[str, int, int]: + return ( + str(slot_mapping.device), + int(slot_mapping.data_ptr()), + int(slot_mapping.numel()), + ) + + def validate_slot_mapping(self, slot_mapping: torch.Tensor) -> None: + latent_cache, _ = self._require_caches() + if slot_mapping.ndim != 1: + raise ValueError( + f"MLA slot_mapping must be 1D, got {tuple(slot_mapping.shape)}." + ) + if slot_mapping.dtype != torch.int32: + raise TypeError( + f"MLA slot_mapping must use torch.int32, got {slot_mapping.dtype}." + ) + if slot_mapping.device != latent_cache.device: + raise ValueError( + "MLA slot_mapping must share the cache device, got " + f"slots={slot_mapping.device} cache={latent_cache.device}." + ) + validate_copy_slot_mapping( + slot_mapping, + cache_slot_count=int(latent_cache.shape[1]), + ) + self._validated_slot_mapping_key = self._slot_mapping_key(slot_mapping) + self._validated_store_calls_remaining = int(latent_cache.shape[0]) + + def store( + self, + layer_idx: int, + slot_mapping: torch.Tensor, + payload: AttentionCacheWrite, + ) -> None: + if not isinstance(payload, MlaLatentWrite): + raise TypeError( + "MlaLatentStorage.store requires MlaLatentWrite, got " + f"{type(payload).__name__}." + ) + destination = self.layer_payload(layer_idx) + slot_mapping_key = self._slot_mapping_key(slot_mapping) + use_prevalidated_mapping = ( + slot_mapping_key == self._validated_slot_mapping_key + and self._validated_store_calls_remaining > 0 + ) + copy_latent_to_cache( + payload.latent, + payload.rope, + slot_mapping, + destination.latent_cache, + destination.rope_cache, + validate_slots=not use_prevalidated_mapping, + ) + if use_prevalidated_mapping: + self._validated_store_calls_remaining -= 1 + if self._validated_store_calls_remaining == 0: + self._validated_slot_mapping_key = None + + def bytes_per_slot_per_layer(self) -> int: + element_size = torch.tensor([], dtype=self.dtype).element_size() + return int((self.kv_lora_rank + self.rope_dim) * element_size) + + def accounting_tensors(self) -> tuple[torch.Tensor, ...]: + return self._require_caches() diff --git a/tests/test_attention_cache_storage.py b/tests/test_attention_cache_storage.py new file mode 100644 index 00000000..b0bd2369 --- /dev/null +++ b/tests/test_attention_cache_storage.py @@ -0,0 +1,424 @@ +from __future__ import annotations + +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + +import numpy as np +import pytest +import torch + +from sparsevllm.config import Config, RuntimeLayout +from sparsevllm.engine.cache_manager import ( + ExplicitKVPayload, + ExplicitKVWrite, + LayerBatchStates, + MlaLatentPayload, + MlaLatentWrite, +) +from sparsevllm.engine.cache_manager.standard import StandardCacheManager +from sparsevllm.engine.cache_manager.storage import ( + CacheLayout, + ExplicitKVStorage, + MlaLatentStorage, + create_attention_cache_storage, +) + + +def _glm_hf_config(**overrides): + values = { + "model_type": "glm4_moe_lite", + "architectures": ["Glm4MoeLiteForCausalLM"], + "torch_dtype": torch.bfloat16, + "max_position_embeddings": 4096, + "hidden_size": 64, + "intermediate_size": 128, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "vocab_size": 128, + "kv_lora_rank": 512, + "qk_rope_head_dim": 64, + "quantization_config": None, + } + values.update(overrides) + return SimpleNamespace(**values) + + +def _glm_config(*, hf_overrides=None, **overrides) -> Config: + kwargs = { + "model": str(Path(__file__).resolve().parents[1]), + "max_model_len": 128, + "max_num_batched_tokens": 64, + "chunk_prefill_size": 64, + "enforce_eager": True, + } + kwargs.update(overrides) + with patch( + "sparsevllm.configs.runtime.AutoConfig.from_pretrained", + return_value=_glm_hf_config(**(hf_overrides or {})), + ): + return Config(**kwargs) + + +def test_glm_config_selects_mla_latent_layout(): + config = _glm_config() + + assert config.attention_cache_layout == CacheLayout.MLA_LATENT.value + + +@pytest.mark.parametrize( + ("override", "message"), + [ + ({"vllm_sparse_method": "snapkv"}, "only vanilla attention"), + ({"enable_prefix_caching": True}, "does not support prefix caching"), + ({"decode_cuda_graph": True}, "does not support decode CUDA Graph"), + ({"enforce_eager": False}, "requires enforce_eager=True"), + ], +) +def test_glm_config_rejects_unsupported_storage_combinations(override, message): + with pytest.raises(NotImplementedError, match=message): + _glm_config(**override) + + +@pytest.mark.parametrize( + ("hf_override", "message"), + [ + ({"torch_dtype": torch.float16}, "requires torch.bfloat16"), + ({"kv_lora_rank": 256}, "requires kv_lora_rank=512"), + ({"qk_rope_head_dim": 32}, "qk_rope_head_dim=64"), + ], +) +def test_glm_config_rejects_unsupported_mla_storage_contract(hf_override, message): + with pytest.raises(NotImplementedError, match=message): + _glm_config(hf_overrides=hf_override) + + +def test_explicit_storage_preserves_legacy_tensor_layout_and_size(): + storage = ExplicitKVStorage( + num_kv_heads=2, + head_dim=8, + dtype=torch.float16, + ) + storage.allocate(num_layers=3, num_slots=5, device=torch.device("cpu")) + + assert storage.layout is CacheLayout.EXPLICIT_KV + assert storage.cache.shape == (2, 3, 5, 2, 8) + assert storage.bytes_per_slot_per_layer() == 2 * 2 * 8 * 2 + assert storage.cache.untyped_storage().nbytes() == 3 * 5 * 2 * 2 * 8 * 2 + payload = storage.layer_payload(1) + assert isinstance(payload, ExplicitKVPayload) + assert payload.k_cache.data_ptr() == storage.cache[0, 1].data_ptr() + assert payload.v_cache.data_ptr() == storage.cache[1, 1].data_ptr() + accounting_tensors = storage.accounting_tensors() + assert len(accounting_tensors) == 1 + assert accounting_tensors[0] is storage.cache + + +def test_mla_storage_uses_576_bf16_values_per_slot_per_layer(): + storage = MlaLatentStorage( + kv_lora_rank=512, + rope_dim=64, + dtype=torch.bfloat16, + ) + storage.allocate(num_layers=2, num_slots=3, device=torch.device("cpu")) + + assert storage.layout is CacheLayout.MLA_LATENT + assert storage.latent_cache is not None + assert storage.rope_cache is not None + assert storage.latent_cache.shape == (2, 3, 1, 512) + assert storage.rope_cache.shape == (2, 3, 1, 64) + assert storage.bytes_per_slot_per_layer() == 576 * 2 + assert sum(t.untyped_storage().nbytes() for t in storage.accounting_tensors()) == ( + 2 * 3 * 576 * 2 + ) + payload = storage.layer_payload(1) + assert isinstance(payload, MlaLatentPayload) + assert payload.latent_cache.data_ptr() == storage.latent_cache[1].data_ptr() + assert payload.rope_cache.data_ptr() == storage.rope_cache[1].data_ptr() + + +def test_storage_factory_uses_configured_layout(): + explicit_config = SimpleNamespace( + attention_cache_layout="explicit_kv", + hf_config=SimpleNamespace(torch_dtype=torch.float16), + ) + mla_config = SimpleNamespace( + attention_cache_layout="mla_latent", + hf_config=SimpleNamespace( + torch_dtype=torch.bfloat16, + kv_lora_rank=512, + qk_rope_head_dim=64, + ), + ) + + assert isinstance( + create_attention_cache_storage( + explicit_config, + num_kv_heads=2, + head_dim=8, + ), + ExplicitKVStorage, + ) + assert isinstance( + create_attention_cache_storage( + mla_config, + num_kv_heads=4, + head_dim=64, + ), + MlaLatentStorage, + ) + + inferred_mla_config = SimpleNamespace( + hf_config=SimpleNamespace( + model_type="glm4_moe_lite", + torch_dtype=torch.bfloat16, + kv_lora_rank=512, + qk_rope_head_dim=64, + ) + ) + assert isinstance( + create_attention_cache_storage( + inferred_mla_config, + num_kv_heads=4, + head_dim=64, + ), + MlaLatentStorage, + ) + + +@pytest.mark.parametrize( + ("storage", "num_layers", "num_slots", "expected_shape"), + [ + ( + ExplicitKVStorage( + num_kv_heads=2, + head_dim=8, + dtype=torch.float16, + ), + 3, + 5, + (2, 3, 5, 2, 8), + ), + ( + MlaLatentStorage( + kv_lora_rank=512, + rope_dim=64, + dtype=torch.bfloat16, + ), + 2, + 7, + ((2, 7, 1, 512), (2, 7, 1, 64)), + ), + ], +) +def test_standard_manager_derives_capacity_and_allocates_through_storage( + storage, + num_layers, + num_slots, + expected_shape, +): + manager = object.__new__(StandardCacheManager) + manager.attention_cache_storage = storage + manager.num_kv_layers = num_layers + manager.device = torch.device("cpu") + manager.config = SimpleNamespace(num_kvcache_slots=-1) + slot_bytes = storage.bytes_per_slot_per_layer() + manager._get_available_slots_info = lambda: ( + num_layers * num_slots * slot_bytes, + slot_bytes, + ) + + manager.allocate_kv_cache() + + assert manager.config.num_kvcache_slots == num_slots + if isinstance(storage, ExplicitKVStorage): + assert storage.cache.shape == expected_shape + assert manager.kv_cache is storage.cache + else: + assert storage.latent_cache is not None + assert storage.rope_cache is not None + assert (storage.latent_cache.shape, storage.rope_cache.shape) == expected_shape + assert manager.kv_cache is None + + +def test_storage_store_payload_types_are_not_interchangeable(): + explicit = ExplicitKVStorage( + num_kv_heads=1, + head_dim=4, + dtype=torch.float16, + ) + explicit.allocate(num_layers=1, num_slots=2, device=torch.device("cpu")) + mla = MlaLatentStorage( + kv_lora_rank=512, + rope_dim=64, + dtype=torch.bfloat16, + ) + mla.allocate(num_layers=1, num_slots=2, device=torch.device("cpu")) + slots = torch.tensor([0], dtype=torch.int32) + + with pytest.raises(TypeError, match="ExplicitKVWrite"): + explicit.store( + 0, + slots, + MlaLatentWrite( + latent=torch.empty(1, 1, 512, dtype=torch.bfloat16), + rope=torch.empty(1, 1, 64, dtype=torch.bfloat16), + ), + ) + with pytest.raises(TypeError, match="MlaLatentWrite"): + mla.store( + 0, + slots, + ExplicitKVWrite( + key=torch.empty(1, 1, 4, dtype=torch.float16), + value=torch.empty(1, 1, 4, dtype=torch.float16), + ), + ) + + +def test_mla_storage_reuses_one_manager_validation_across_layers(): + storage = MlaLatentStorage( + kv_lora_rank=512, + rope_dim=64, + dtype=torch.bfloat16, + ) + storage.allocate(num_layers=2, num_slots=2, device=torch.device("cpu")) + slots = torch.tensor([0], dtype=torch.int32) + write = MlaLatentWrite( + latent=torch.empty(1, 1, 512, dtype=torch.bfloat16), + rope=torch.empty(1, 1, 64, dtype=torch.bfloat16), + ) + storage.validate_slot_mapping(slots) + + with patch( + "sparsevllm.engine.cache_manager.storage.mla_latent.copy_latent_to_cache" + ) as copy: + storage.store(0, slots, write) + storage.store(1, slots, write) + storage.store(0, slots, write) + + assert [call.kwargs["validate_slots"] for call in copy.call_args_list] == [ + False, + False, + True, + ] + + +def test_standard_manager_delegates_payload_store_and_compute_view(): + storage = MlaLatentStorage( + kv_lora_rank=512, + rope_dim=64, + dtype=torch.bfloat16, + ) + storage.allocate(num_layers=2, num_slots=4, device=torch.device("cpu")) + manager = object.__new__(StandardCacheManager) + manager.attention_cache_storage = storage + manager.runtime_layout = RuntimeLayout.dense(2) + manager.layer_batch_state = LayerBatchStates( + slot_mapping=torch.tensor([1], dtype=torch.int32) + ) + write = MlaLatentWrite( + latent=torch.empty(1, 1, 512, dtype=torch.bfloat16), + rope=torch.empty(1, 1, 64, dtype=torch.bfloat16), + ) + + with patch.object(storage, "store") as store: + returned_slots = manager.store_attention_payload(1, write) + store.assert_called_once_with(1, manager.layer_batch_state.slot_mapping, write) + assert returned_slots is manager.layer_batch_state.slot_mapping + + active_slots = torch.tensor([[0, 1]], dtype=torch.int32) + req_indices = torch.tensor([0], dtype=torch.int32) + context_lens = torch.tensor([2], dtype=torch.int32) + payload, actual_slots, actual_rows, actual_lens = manager.get_layer_compute_payload( + 1, + active_slots, + req_indices, + context_lens, + ) + assert isinstance(payload, MlaLatentPayload) + assert actual_slots is active_slots + assert actual_rows is req_indices + assert actual_lens is context_lens + + +def test_standard_manager_accounts_storage_tensors_explicitly(): + storage = MlaLatentStorage( + kv_lora_rank=512, + rope_dim=64, + dtype=torch.bfloat16, + ) + storage.allocate(num_layers=2, num_slots=3, device=torch.device("cpu")) + manager = object.__new__(StandardCacheManager) + manager.attention_cache_storage = storage + manager.kv_cache = None + manager.config = SimpleNamespace( + num_kvcache_slots=3, + max_num_seqs_in_gpu=1, + memory_expected_savings=None, + ) + manager.hf_config = SimpleNamespace(torch_dtype=torch.bfloat16) + manager.num_layers = 2 + manager.num_kv_layers = 2 + manager.num_kv_heads = 4 + manager.head_dim = 64 + manager.row_seq_lens = np.array([2], dtype=np.int32) + + accounting = manager.memory_accounting() + + assert accounting["kv_or_latent_tensor_bytes"] == 2 * 3 * 576 * 2 + assert accounting["logical_live_kv_bytes"] == 2 * 2 * 576 * 2 + assert accounting["tensor_count"] == 2 + assert {item["path"] for item in accounting["tensors"]} == { + "attention_cache_storage.mla_latent.0_cache", + "attention_cache_storage.mla_latent.1_cache", + } + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_mla_storage_store_skips_padding_and_overwrites_reused_slot(): + device = torch.device("cuda") + storage = MlaLatentStorage( + kv_lora_rank=512, + rope_dim=64, + dtype=torch.bfloat16, + ) + storage.allocate(num_layers=1, num_slots=4, device=device) + assert storage.latent_cache is not None + assert storage.rope_cache is not None + storage.latent_cache.fill_(-7) + storage.rope_cache.fill_(-7) + + latent = torch.stack( + [torch.full((1, 512), value, dtype=torch.bfloat16, device=device) for value in (1, 2, 3)] + ) + rope = torch.stack( + [torch.full((1, 64), value, dtype=torch.bfloat16, device=device) for value in (4, 5, 6)] + ) + slot_mapping = torch.tensor([1, -1, 3], dtype=torch.int32, device=device) + storage.validate_slot_mapping(slot_mapping) + storage.store( + 0, + slot_mapping, + MlaLatentWrite(latent=latent, rope=rope), + ) + + assert torch.equal(storage.latent_cache[0, 0], torch.full_like(storage.latent_cache[0, 0], -7)) + assert torch.equal(storage.latent_cache[0, 1], latent[0]) + assert torch.equal(storage.latent_cache[0, 2], torch.full_like(storage.latent_cache[0, 2], -7)) + assert torch.equal(storage.latent_cache[0, 3], latent[2]) + assert torch.equal(storage.rope_cache[0, 1], rope[0]) + assert torch.equal(storage.rope_cache[0, 3], rope[2]) + + replacement = MlaLatentWrite( + latent=torch.full((1, 1, 512), 9, dtype=torch.bfloat16, device=device), + rope=torch.full((1, 1, 64), 10, dtype=torch.bfloat16, device=device), + ) + storage.store( + 0, + torch.tensor([1], dtype=torch.int32, device=device), + replacement, + ) + assert torch.equal(storage.latent_cache[0, 1], replacement.latent[0]) + assert torch.equal(storage.rope_cache[0, 1], replacement.rope[0]) From a27a81f94742e7b9e891d5d7bf5245e9af0fc9bc Mon Sep 17 00:00:00 2001 From: Codex Date: Tue, 4 Aug 2026 22:36:29 +0800 Subject: [PATCH 04/48] feat: add mla attention provider --- src/sparsevllm/operators/mla_attention.py | 314 ++++++++++++++++++ .../triton_kernel/mla/decode_schedule.py | 9 +- tests/test_mla_attention_operator.py | 314 ++++++++++++++++++ tests/test_mla_kernels.py | 5 + 4 files changed, 641 insertions(+), 1 deletion(-) create mode 100644 src/sparsevllm/operators/mla_attention.py create mode 100644 tests/test_mla_attention_operator.py diff --git a/src/sparsevllm/operators/mla_attention.py b/src/sparsevllm/operators/mla_attention.py new file mode 100644 index 00000000..2f9a6806 --- /dev/null +++ b/src/sparsevllm/operators/mla_attention.py @@ -0,0 +1,314 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +import sparsevllm.platforms as platforms +from sparsevllm.engine.cache_manager.base import ( + DecodeComputeView, + MlaLatentPayload, +) +from sparsevllm.operators.registry import ( + OpRegistry, + OpResolver, + SupportResult, +) +from sparsevllm.platforms.interface import DeviceCaps, PlatformEnum +from sparsevllm.triton_kernel.mla import ( + DEFAULT_GLM_MLA_DECODE_CONFIG, + MlaDecodeLaunchConfig, + allocate_mla_decode_workspace, + run_mla_decode, +) + + +_GLM_MLA_NUM_Q_HEADS = 20 +_GLM_MLA_KV_LORA_RANK = 512 +_GLM_MLA_ROPE_DIM = 64 +_GLM_MLA_QK_HEAD_DIM = 256 +_GLM_MLA_VALUE_HEAD_DIM = 256 +_VALIDATED_H100_NAME = "NVIDIA H100 80GB HBM3" + + +@dataclass(frozen=True, slots=True) +class MlaAttentionOpSpec: + """Construction-time contract for one MLA attention implementation.""" + + num_q_heads: int + kv_lora_rank: int + rope_dim: int + qk_head_dim: int + value_head_dim: int + activation_dtype: torch.dtype + cache_dtype: torch.dtype + tp_size: int + cuda_graph: bool + + def __post_init__(self) -> None: + dimensions = { + "num_q_heads": self.num_q_heads, + "kv_lora_rank": self.kv_lora_rank, + "rope_dim": self.rope_dim, + "qk_head_dim": self.qk_head_dim, + "value_head_dim": self.value_head_dim, + "tp_size": self.tp_size, + } + for name, value in dimensions.items(): + if int(value) <= 0: + raise ValueError(f"MLA {name} must be positive, got {value}.") + if self.num_q_heads % self.tp_size: + raise ValueError( + "MLA query heads must be divisible by tensor parallel size: " + f"heads={self.num_q_heads} tp_size={self.tp_size}." + ) + + @property + def local_q_heads(self) -> int: + return int(self.num_q_heads // self.tp_size) + + @property + def softmax_scale(self) -> float: + return float(self.qk_head_dim**-0.5) + + +class MlaAttentionProvider: + name = "" + priority = 0 + + def run( + self, + q_nope_absorbed: torch.Tensor, + q_rope: torch.Tensor, + view: DecodeComputeView, + output: torch.Tensor, + ) -> torch.Tensor: + raise NotImplementedError + + +MLA_ATTENTION_REGISTRY: OpRegistry[ + MlaAttentionOpSpec, + MlaAttentionProvider, +] = OpRegistry("MLA attention") + + +@MLA_ATTENTION_REGISTRY.register +class MlaTritonProvider(MlaAttentionProvider): + """H100 provider bound once with caller-independent decode workspace.""" + + name = "triton_h100" + priority = 100 + + def __init__( + self, + *, + op_spec: MlaAttentionOpSpec, + device: torch.device | str, + max_batch_size: int, + launch_config: MlaDecodeLaunchConfig = DEFAULT_GLM_MLA_DECODE_CONFIG, + ) -> None: + self.spec = op_spec + requested_device = torch.device(device) + self.max_batch_size = int(max_batch_size) + if self.max_batch_size <= 0: + raise ValueError( + "MLA max_batch_size must be positive, got " + f"{self.max_batch_size}." + ) + self.launch_config = launch_config + self.workspace = allocate_mla_decode_workspace( + batch_size=self.max_batch_size, + head_count=self.spec.local_q_heads, + device=requested_device, + config=self.launch_config, + ) + self.device = self.workspace.block_size.device + + @classmethod + def supports( + cls, + spec: MlaAttentionOpSpec, + caps: DeviceCaps, + ) -> SupportResult: + if caps.platform != PlatformEnum.CUDA or caps.compute_capability != (9, 0): + return SupportResult.no( + f"requires CUDA SM90, got {caps.platform.name} " + f"{caps.compute_capability}" + ) + if caps.device_name != _VALIDATED_H100_NAME: + return SupportResult.no( + "requires validated NVIDIA H100 80GB HBM3 hardware, got " + f"{caps.device_name}" + ) + if not caps.supports_triton: + return SupportResult.no("platform does not support Triton") + if not caps.supports_bfloat16: + return SupportResult.no("device does not support BF16") + if spec.activation_dtype != torch.bfloat16: + return SupportResult.no( + f"requires BF16 activations, got {spec.activation_dtype}" + ) + if spec.cache_dtype != torch.bfloat16: + return SupportResult.no( + f"requires BF16 cache storage, got {spec.cache_dtype}" + ) + expected_shape = ( + _GLM_MLA_NUM_Q_HEADS, + _GLM_MLA_KV_LORA_RANK, + _GLM_MLA_ROPE_DIM, + _GLM_MLA_QK_HEAD_DIM, + _GLM_MLA_VALUE_HEAD_DIM, + ) + actual_shape = ( + spec.num_q_heads, + spec.kv_lora_rank, + spec.rope_dim, + spec.qk_head_dim, + spec.value_head_dim, + ) + if actual_shape != expected_shape: + return SupportResult.no( + f"requires GLM MLA shape {expected_shape}, got {actual_shape}" + ) + if spec.tp_size not in {1, 2, 4}: + return SupportResult.no( + f"requires tensor parallel size 1, 2, or 4, got {spec.tp_size}" + ) + if spec.cuda_graph: + return SupportResult.no("CUDA Graph execution is not validated") + return SupportResult.yes() + + def _validate_run_inputs( + self, + q_nope_absorbed: torch.Tensor, + q_rope: torch.Tensor, + view: DecodeComputeView, + output: torch.Tensor, + ) -> MlaLatentPayload: + if not isinstance(view, DecodeComputeView): + raise TypeError( + "MlaTritonProvider.run requires DecodeComputeView, got " + f"{type(view).__name__}." + ) + if not isinstance(view.payload, MlaLatentPayload): + raise TypeError( + "MLA decode requires MlaLatentPayload, got " + f"{type(view.payload).__name__}." + ) + if q_nope_absorbed.ndim != 3: + raise ValueError( + "q_nope_absorbed must have shape [batch, local_heads, 512], " + f"got {tuple(q_nope_absorbed.shape)}." + ) + expected_query_shape = ( + int(q_nope_absorbed.shape[0]), + self.spec.local_q_heads, + self.spec.kv_lora_rank, + ) + if tuple(q_nope_absorbed.shape) != expected_query_shape: + raise ValueError( + "q_nope_absorbed must have shape " + f"{expected_query_shape}, got {tuple(q_nope_absorbed.shape)}." + ) + expected_rope_shape = ( + expected_query_shape[0], + expected_query_shape[1], + self.spec.rope_dim, + ) + if tuple(q_rope.shape) != expected_rope_shape: + raise ValueError( + f"q_rope must have shape {expected_rope_shape}, got " + f"{tuple(q_rope.shape)}." + ) + if output.shape != q_nope_absorbed.shape: + raise ValueError( + f"output must have shape {tuple(q_nope_absorbed.shape)}, got " + f"{tuple(output.shape)}." + ) + if expected_query_shape[0] > self.max_batch_size: + raise ValueError( + "MLA decode batch exceeds the bound workspace: " + f"batch={expected_query_shape[0]} max_batch_size=" + f"{self.max_batch_size}." + ) + tensors = { + "q_nope_absorbed": q_nope_absorbed, + "q_rope": q_rope, + "output": output, + "latent_cache": view.payload.latent_cache, + "rope_cache": view.payload.rope_cache, + } + for name, tensor in tensors.items(): + if tensor.device != self.device: + raise ValueError( + f"{name} is on {tensor.device}, expected {self.device}." + ) + expected_dtype = ( + self.spec.cache_dtype + if name in {"latent_cache", "rope_cache"} + else self.spec.activation_dtype + ) + if tensor.dtype != expected_dtype: + raise TypeError( + f"{name} must use {expected_dtype}, got {tensor.dtype}." + ) + return view.payload + + @torch.no_grad() + def run( + self, + q_nope_absorbed: torch.Tensor, + q_rope: torch.Tensor, + view: DecodeComputeView, + output: torch.Tensor, + ) -> torch.Tensor: + payload = self._validate_run_inputs( + q_nope_absorbed, + q_rope, + view, + output, + ) + return run_mla_decode( + q_nope_absorbed, + q_rope, + payload.latent_cache, + payload.rope_cache, + view.meta.active_slots, + view.meta.req_indices, + view.meta.context_lens, + output, + self.workspace, + softmax_scale=self.spec.softmax_scale, + config=self.launch_config, + ) + + +def resolve_mla_attention_provider( + spec: MlaAttentionOpSpec, + *, + device: torch.device | str, + max_batch_size: int, + launch_config: MlaDecodeLaunchConfig = DEFAULT_GLM_MLA_DECODE_CONFIG, +) -> MlaAttentionProvider: + """Resolve and bind an MLA provider during model construction.""" + + device = torch.device(device) + device_index = 0 if device.index is None else int(device.index) + caps = platforms.current_platform.get_device_caps(device_index) + return OpResolver(MLA_ATTENTION_REGISTRY).resolve( + spec, + caps, + op_spec=spec, + device=device, + max_batch_size=max_batch_size, + launch_config=launch_config, + ).provider + + +__all__ = [ + "MLA_ATTENTION_REGISTRY", + "MlaAttentionOpSpec", + "MlaAttentionProvider", + "MlaTritonProvider", + "resolve_mla_attention_provider", +] diff --git a/src/sparsevllm/triton_kernel/mla/decode_schedule.py b/src/sparsevllm/triton_kernel/mla/decode_schedule.py index 09d563f9..8790e38a 100644 --- a/src/sparsevllm/triton_kernel/mla/decode_schedule.py +++ b/src/sparsevllm/triton_kernel/mla/decode_schedule.py @@ -8,6 +8,7 @@ from __future__ import annotations +import math from dataclasses import dataclass import torch @@ -326,6 +327,7 @@ def run_mla_decode( output: torch.Tensor, workspace: MlaDecodeWorkspace, *, + softmax_scale: float, config: MlaDecodeLaunchConfig = DEFAULT_GLM_MLA_DECODE_CONFIG, validate_metadata: bool = True, ) -> torch.Tensor: @@ -358,6 +360,11 @@ def run_mla_decode( "MLA decode workspace is too small: " f"heads={workspace.mid_output.shape[0]}, required={head_count}" ) + softmax_scale = float(softmax_scale) + if not math.isfinite(softmax_scale) or softmax_scale <= 0: + raise ValueError( + f"softmax_scale must be finite and positive, got {softmax_scale}." + ) if validate_metadata: validate_mla_decode_metadata( active_slots, @@ -378,7 +385,7 @@ def run_mla_decode( workspace.block_size, workspace.mid_output, workspace.mid_logsumexp, - softmax_scale=GLM_MLA_SOFTMAX_SCALE, + softmax_scale=softmax_scale, program_count=config.program_count, block_q_heads=config.block_q_heads, block_n=config.block_n, diff --git a/tests/test_mla_attention_operator.py b/tests/test_mla_attention_operator.py new file mode 100644 index 00000000..db39bdd3 --- /dev/null +++ b/tests/test_mla_attention_operator.py @@ -0,0 +1,314 @@ +from __future__ import annotations + +from unittest.mock import patch + +import pytest +import torch + +from sparsevllm.engine.cache_manager import ( + AttentionViewMeta, + DecodeComputeView, + ExplicitKVPayload, + MlaLatentPayload, +) +from sparsevllm.operators.mla_attention import ( + MLA_ATTENTION_REGISTRY, + MlaAttentionOpSpec, + MlaTritonProvider, +) +from sparsevllm.operators.registry import OpResolver +from sparsevllm.platforms import DeviceCaps, PlatformEnum +from sparsevllm.triton_kernel.mla import MlaDecodeWorkspace + + +def _spec(**overrides) -> MlaAttentionOpSpec: + values = { + "num_q_heads": 20, + "kv_lora_rank": 512, + "rope_dim": 64, + "qk_head_dim": 256, + "value_head_dim": 256, + "activation_dtype": torch.bfloat16, + "cache_dtype": torch.bfloat16, + "tp_size": 4, + "cuda_graph": False, + } + values.update(overrides) + return MlaAttentionOpSpec(**values) + + +def _h100_caps(**overrides) -> DeviceCaps: + values = { + "platform": PlatformEnum.CUDA, + "device_type": "cuda", + "device_index": 0, + "device_name": "NVIDIA H100 80GB HBM3", + "compute_capability": (9, 0), + "runtime_version": "12.9", + "supports_graph_capture": True, + "supports_torch_compile": True, + "supports_triton": True, + "supports_pin_memory": True, + "supports_bfloat16": True, + "supports_native_fp8": True, + } + values.update(overrides) + return DeviceCaps(**values) + + +def _cpu_workspace(batch_size: int, head_count: int) -> MlaDecodeWorkspace: + return MlaDecodeWorkspace( + block_size=torch.empty(1, dtype=torch.int32), + batch_start_indices=torch.empty(batch_size, dtype=torch.int32), + mid_output=torch.empty(head_count, 1, 512, dtype=torch.float32), + mid_logsumexp=torch.empty(head_count, 1, dtype=torch.float32), + ) + + +@pytest.mark.parametrize( + "overrides", + [ + {"num_q_heads": 0}, + {"kv_lora_rank": 0}, + {"rope_dim": -1}, + {"qk_head_dim": 0}, + {"value_head_dim": 0}, + {"tp_size": 0}, + {"num_q_heads": 20, "tp_size": 3}, + ], +) +def test_mla_attention_spec_rejects_invalid_dimensions(overrides) -> None: + with pytest.raises(ValueError): + _spec(**overrides) + + +def test_mla_attention_scale_uses_qk_head_dimension() -> None: + spec = _spec() + + assert spec.softmax_scale == pytest.approx(256**-0.5) + assert spec.softmax_scale != pytest.approx((512 + 64) ** -0.5) + + +def test_mla_resolver_selects_h100_triton_provider() -> None: + spec = _spec() + workspace = _cpu_workspace(batch_size=8, head_count=spec.local_q_heads) + + with patch( + "sparsevllm.operators.mla_attention.allocate_mla_decode_workspace", + return_value=workspace, + ) as allocate: + resolved = OpResolver(MLA_ATTENTION_REGISTRY).resolve( + spec, + _h100_caps(), + op_spec=spec, + device="cpu", + max_batch_size=8, + ) + + assert isinstance(resolved.provider, MlaTritonProvider) + assert resolved.rejected == () + allocate.assert_called_once_with( + batch_size=8, + head_count=5, + device=torch.device("cpu"), + config=resolved.provider.launch_config, + ) + + +@pytest.mark.parametrize( + ("spec_overrides", "caps_overrides", "reason"), + [ + ({}, {"platform": PlatformEnum.CPU}, "requires CUDA SM90"), + ({}, {"compute_capability": (8, 0)}, "requires CUDA SM90"), + ({}, {"device_name": "NVIDIA H100 PCIe"}, "validated"), + ({}, {"supports_triton": False}, "does not support Triton"), + ({}, {"supports_bfloat16": False}, "does not support BF16"), + ({"activation_dtype": torch.float16}, {}, "BF16 activations"), + ({"cache_dtype": torch.float16}, {}, "BF16 cache"), + ({"kv_lora_rank": 256}, {}, "GLM MLA shape"), + ({"tp_size": 5}, {}, "tensor parallel size"), + ({"cuda_graph": True}, {}, "CUDA Graph"), + ], +) +def test_mla_resolver_rejects_unvalidated_contracts( + spec_overrides, + caps_overrides, + reason, +) -> None: + with pytest.raises(RuntimeError, match=reason): + OpResolver(MLA_ATTENTION_REGISTRY).resolve( + _spec(**spec_overrides), + _h100_caps(**caps_overrides), + ) + + +def test_mla_provider_rejects_explicit_kv_before_kernel() -> None: + spec = _spec(tp_size=1) + workspace = _cpu_workspace(batch_size=1, head_count=20) + with patch( + "sparsevllm.operators.mla_attention.allocate_mla_decode_workspace", + return_value=workspace, + ): + provider = MlaTritonProvider( + op_spec=spec, + device="cpu", + max_batch_size=1, + ) + view = DecodeComputeView( + meta=AttentionViewMeta( + active_slots=torch.tensor([[0]], dtype=torch.int32), + req_indices=torch.tensor([0], dtype=torch.int32), + context_lens=torch.tensor([1], dtype=torch.int32), + ), + payload=ExplicitKVPayload( + k_cache=torch.empty(1, 1, 256, dtype=torch.bfloat16), + v_cache=torch.empty(1, 1, 256, dtype=torch.bfloat16), + ), + ) + + with ( + patch("sparsevllm.operators.mla_attention.run_mla_decode") as kernel, + pytest.raises(TypeError, match="MlaLatentPayload"), + ): + provider.run( + torch.empty(1, 20, 512, dtype=torch.bfloat16), + torch.empty(1, 20, 64, dtype=torch.bfloat16), + view, + torch.empty(1, 20, 512, dtype=torch.bfloat16), + ) + kernel.assert_not_called() + + +def test_mla_provider_run_does_not_resolve_or_allocate() -> None: + spec = _spec(tp_size=4) + workspace = _cpu_workspace(batch_size=2, head_count=5) + with patch( + "sparsevllm.operators.mla_attention.allocate_mla_decode_workspace", + return_value=workspace, + ): + provider = MlaTritonProvider( + op_spec=spec, + device="cpu", + max_batch_size=2, + ) + payload = MlaLatentPayload( + latent_cache=torch.empty(4, 1, 512, dtype=torch.bfloat16), + rope_cache=torch.empty(4, 1, 64, dtype=torch.bfloat16), + ) + view = DecodeComputeView( + meta=AttentionViewMeta( + active_slots=torch.tensor([[0, 1], [2, 3]], dtype=torch.int32), + req_indices=torch.tensor([0, -1], dtype=torch.int32), + context_lens=torch.tensor([2, 0], dtype=torch.int32), + ), + payload=payload, + ) + q_nope_absorbed = torch.empty(2, 5, 512, dtype=torch.bfloat16) + q_rope = torch.empty(2, 5, 64, dtype=torch.bfloat16) + output = torch.empty_like(q_nope_absorbed) + + with ( + patch.object(OpResolver, "resolve") as resolve, + patch( + "sparsevllm.operators.mla_attention.allocate_mla_decode_workspace" + ) as allocate, + patch( + "sparsevllm.operators.mla_attention.run_mla_decode", + return_value=output, + ) as kernel, + ): + actual = provider.run(q_nope_absorbed, q_rope, view, output) + + assert actual is output + resolve.assert_not_called() + allocate.assert_not_called() + kernel.assert_called_once_with( + q_nope_absorbed, + q_rope, + payload.latent_cache, + payload.rope_cache, + view.meta.active_slots, + view.meta.req_indices, + view.meta.context_lens, + output, + workspace, + softmax_scale=spec.softmax_scale, + config=provider.launch_config, + ) + + +def test_mla_provider_rejects_batch_larger_than_workspace() -> None: + spec = _spec(tp_size=4) + workspace = _cpu_workspace(batch_size=1, head_count=5) + with patch( + "sparsevllm.operators.mla_attention.allocate_mla_decode_workspace", + return_value=workspace, + ): + provider = MlaTritonProvider( + op_spec=spec, + device="cpu", + max_batch_size=1, + ) + view = DecodeComputeView( + meta=AttentionViewMeta( + active_slots=torch.tensor([[0], [1]], dtype=torch.int32), + req_indices=torch.tensor([0, 1], dtype=torch.int32), + context_lens=torch.tensor([1, 1], dtype=torch.int32), + ), + payload=MlaLatentPayload( + latent_cache=torch.empty(2, 1, 512, dtype=torch.bfloat16), + rope_cache=torch.empty(2, 1, 64, dtype=torch.bfloat16), + ), + ) + + with pytest.raises(ValueError, match="exceeds the bound workspace"): + provider.run( + torch.empty(2, 5, 512, dtype=torch.bfloat16), + torch.empty(2, 5, 64, dtype=torch.bfloat16), + view, + torch.empty(2, 5, 512, dtype=torch.bfloat16), + ) + + +@pytest.mark.skipif( + not torch.cuda.is_available(), + reason="CUDA is required for the MLA provider integration test", +) +def test_mla_provider_runs_static_padded_batch() -> None: + spec = _spec(tp_size=4) + provider = MlaTritonProvider( + op_spec=spec, + device="cuda", + max_batch_size=2, + ) + q_nope_absorbed = torch.randn( + 2, + 5, + 512, + dtype=torch.bfloat16, + device="cuda", + ) + q_rope = torch.randn(2, 5, 64, dtype=torch.bfloat16, device="cuda") + payload = MlaLatentPayload( + latent_cache=torch.randn(4, 1, 512, dtype=torch.bfloat16, device="cuda"), + rope_cache=torch.randn(4, 1, 64, dtype=torch.bfloat16, device="cuda"), + ) + view = DecodeComputeView( + meta=AttentionViewMeta( + active_slots=torch.tensor( + [[0, 2], [1, 3]], + dtype=torch.int32, + device="cuda", + ), + req_indices=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), + context_lens=torch.tensor([2, 0], dtype=torch.int32, device="cuda"), + ), + payload=payload, + ) + output = torch.empty_like(q_nope_absorbed) + + provider.run(q_nope_absorbed, q_rope, view, output) + torch.cuda.synchronize() + + torch.testing.assert_close(output[1], torch.zeros_like(output[1])) + assert bool(torch.isfinite(output).all().item()) diff --git a/tests/test_mla_kernels.py b/tests/test_mla_kernels.py index 47c5b7ae..e59fd278 100644 --- a/tests/test_mla_kernels.py +++ b/tests/test_mla_kernels.py @@ -326,6 +326,7 @@ def test_mla_decode_matches_torch_matrix( context_lens, output, workspace, + softmax_scale=GLM_MLA_SOFTMAX_SCALE, ) torch.cuda.synchronize() expected = _torch_mla_decode(*case) @@ -380,6 +381,7 @@ def test_mla_decode_zeroes_padded_rows() -> None: context_lens, output, workspace, + softmax_scale=GLM_MLA_SOFTMAX_SCALE, ) torch.cuda.synchronize() @@ -423,6 +425,7 @@ def test_mla_decode_rejects_duplicate_active_slots() -> None: context_lens, output, workspace, + softmax_scale=GLM_MLA_SOFTMAX_SCALE, ) @@ -565,6 +568,7 @@ def test_decode_rejects_small_workspace_and_non_bf16_input() -> None: context_lens, output, small_workspace, + softmax_scale=GLM_MLA_SOFTMAX_SCALE, ) with pytest.raises(TypeError, match="q_latent"): run_mla_decode( @@ -577,4 +581,5 @@ def test_decode_rejects_small_workspace_and_non_bf16_input() -> None: context_lens, output, workspace, + softmax_scale=GLM_MLA_SOFTMAX_SCALE, ) From 1ff41eb82d5972f48ad1667b76f796d7c5c8faf8 Mon Sep 17 00:00:00 2001 From: Codex Date: Tue, 4 Aug 2026 22:55:11 +0800 Subject: [PATCH 05/48] feat: add mla prefill and decode --- src/sparsevllm/configs/runtime.py | 1 + src/sparsevllm/configs/scheduling.py | 6 + src/sparsevllm/layers/mla_attention.py | 583 ++++++++++++++++++++++ src/sparsevllm/operators/mla_attention.py | 41 ++ src/sparsevllm/utils/context.py | 1 + tests/test_attention_cache_storage.py | 6 + tests/test_mla_attention_layer.py | 367 ++++++++++++++ tests/test_mla_attention_operator.py | 37 +- 8 files changed, 1040 insertions(+), 2 deletions(-) create mode 100644 src/sparsevllm/layers/mla_attention.py create mode 100644 tests/test_mla_attention_layer.py diff --git a/src/sparsevllm/configs/runtime.py b/src/sparsevllm/configs/runtime.py index 0c1a1467..c4d1722c 100644 --- a/src/sparsevllm/configs/runtime.py +++ b/src/sparsevllm/configs/runtime.py @@ -61,6 +61,7 @@ class Config( chunk_prefill_size: int | None = None long_prefill_offload_threshold: int = 64 * 1024 mlp_chunk_size: int = 16384 + mla_prefill_workspace_bytes: int = 2 * 1024**3 prefill_schedule_policy: str = PREFILL_POLICY_AUTO gpu_memory_utilization: float = 0.8 device_memory_utilization: float | None = None diff --git a/src/sparsevllm/configs/scheduling.py b/src/sparsevllm/configs/scheduling.py index d575079e..f2edbcc6 100644 --- a/src/sparsevllm/configs/scheduling.py +++ b/src/sparsevllm/configs/scheduling.py @@ -111,3 +111,9 @@ def normalize_scheduling(config) -> None: if int(config.mlp_chunk_size) <= 0: raise ValueError(f"mlp_chunk_size must be > 0, got {config.mlp_chunk_size}.") config.mlp_chunk_size = int(config.mlp_chunk_size) + config.mla_prefill_workspace_bytes = int(config.mla_prefill_workspace_bytes) + if config.mla_prefill_workspace_bytes <= 0: + raise ValueError( + "mla_prefill_workspace_bytes must be > 0, got " + f"{config.mla_prefill_workspace_bytes}." + ) diff --git a/src/sparsevllm/layers/mla_attention.py b/src/sparsevllm/layers/mla_attention.py new file mode 100644 index 00000000..3a4b0ee0 --- /dev/null +++ b/src/sparsevllm/layers/mla_attention.py @@ -0,0 +1,583 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from sparsevllm.engine.cache_manager.base import ( + AttentionViewMeta, + DecodeComputeView, + ExplicitKVPayload, + MlaLatentPayload, + PrefillComputeView, +) +from sparsevllm.layers.attention_backend import TritonAttentionBackend +from sparsevllm.operators.mla_attention import ( + MlaAttentionOpSpec, + MlaAttentionProvider, + resolve_mla_attention_provider, +) +from sparsevllm.triton_kernel.mla import ( + gather_latent_history, + validate_gather_metadata, +) +from sparsevllm.utils.context import get_context + + +@dataclass(frozen=True, slots=True) +class MlaPrefillHistory: + """Gathered full history and its packed logical coordinates.""" + + gathered_latent: torch.Tensor + gathered_rope: torch.Tensor + packed_offsets: torch.Tensor + packed_slots: torch.Tensor + local_req_indices: torch.Tensor + context_lens: torch.Tensor + max_context_len: int + required_workspace_bytes: int + + @property + def visible_tokens(self) -> int: + return int(self.gathered_latent.shape[0]) + + +@dataclass(frozen=True, slots=True) +class MlaPrefillWorkset: + """Full-history MLA buffers ready for ordinary 256-wide attention.""" + + history: MlaPrefillHistory + expanded_k: torch.Tensor + expanded_v: torch.Tensor + + +@dataclass(frozen=True, slots=True) +class _MlaPrefillPlan: + """Step-local packing metadata shared by every MLA layer.""" + + validation_scope: object + source_active_slots: torch.Tensor + source_req_indices: torch.Tensor + source_context_lens: torch.Tensor + source_max_context_len: int | None + cache_slot_count: int + packed_offsets: torch.Tensor + packed_slots: torch.Tensor + local_req_indices: torch.Tensor + total_visible_tokens: int + max_context_len: int + required_workspace_bytes: int + + def matches( + self, + validation_scope: object, + meta: AttentionViewMeta, + cache_slot_count: int, + ) -> bool: + return ( + self.validation_scope is validation_scope + and self.source_active_slots is meta.active_slots + and self.source_req_indices is meta.req_indices + and self.source_context_lens is meta.context_lens + and self.source_max_context_len == meta.max_context_len + and self.cache_slot_count == int(cache_slot_count) + ) + + +def estimate_mla_prefill_workspace_bytes( + *, + total_visible_tokens: int, + batch_size: int, + max_context_len: int, + local_q_heads: int, + activation_dtype: torch.dtype, + cache_dtype: torch.dtype, +) -> int: + """Conservatively bound full-history gather and K/V expansion storage.""" + + values = { + "total_visible_tokens": total_visible_tokens, + "batch_size": batch_size, + "max_context_len": max_context_len, + "local_q_heads": local_q_heads, + } + for name, value in values.items(): + if int(value) < 0: + raise ValueError(f"{name} must be non-negative, got {value}.") + if batch_size == 0 or local_q_heads == 0: + raise ValueError("batch_size and local_q_heads must be positive.") + + cache_element_size = torch.empty((), dtype=cache_dtype).element_size() + activation_element_size = torch.empty( + (), + dtype=activation_dtype, + ).element_size() + gathered_values = int(total_visible_tokens) * (512 + 64) + # kv_b projection materializes 192 K-noPE + 256 V values per head. The + # final K concatenation additionally materializes 192+64 values per head; + # expanded_v may remain a view of the projection output. + expanded_values = ( + int(total_visible_tokens) + * int(local_q_heads) + * ((192 + 256) + 256) + ) + metadata_values = ( + int(batch_size) * int(max_context_len) + + 2 * int(batch_size) + + int(max_context_len) + ) + return int( + gathered_values * cache_element_size + + expanded_values * activation_element_size + + metadata_values * torch.empty((), dtype=torch.int32).element_size() + ) + + +class MLAAttention: + """Semantic MLA execution over tagged cache views. + + Model code owns projection weights, query absorption, and V reconstruction. + This object owns provider binding, decode workspace, full-history gathering, + and reuse of the existing 256-wide prefill attention backend. + """ + + def __init__( + self, + *, + spec: MlaAttentionOpSpec, + provider: MlaAttentionProvider, + prefill_workspace_bytes: int, + ) -> None: + self.spec = spec + self.provider = provider + provider_spec = getattr(provider, "spec", None) + if provider_spec is not None and provider_spec != spec: + raise ValueError( + "MLA semantic layer and provider specs must match: " + f"layer={spec!r} provider={provider_spec!r}." + ) + self.max_batch_size = int(getattr(provider, "max_batch_size", 0)) + if self.max_batch_size <= 0: + raise ValueError("MLA provider must expose a positive max_batch_size.") + self.prefill_workspace_bytes = int(prefill_workspace_bytes) + if self.prefill_workspace_bytes <= 0: + raise ValueError( + "MLA prefill_workspace_bytes must be positive, got " + f"{self.prefill_workspace_bytes}." + ) + if self.spec.qk_head_dim != self.spec.value_head_dim: + raise ValueError( + "The existing prefill backend requires equal QK/value widths, " + f"got {self.spec.qk_head_dim}/{self.spec.value_head_dim}." + ) + self.prefill_backend = TritonAttentionBackend() + self._prefill_plan: _MlaPrefillPlan | None = None + + @classmethod + def bind( + cls, + *, + spec: MlaAttentionOpSpec, + device: torch.device | str, + max_batch_size: int, + prefill_workspace_bytes: int, + ) -> "MLAAttention": + provider = resolve_mla_attention_provider( + spec, + device=device, + max_batch_size=max_batch_size, + ) + return cls( + spec=spec, + provider=provider, + prefill_workspace_bytes=prefill_workspace_bytes, + ) + + @property + def device(self) -> torch.device: + return torch.device(getattr(self.provider, "device")) + + def _require_mla_payload( + self, + view: PrefillComputeView | DecodeComputeView, + *, + operation: str, + ) -> MlaLatentPayload: + payload = view.payload + if not isinstance(payload, MlaLatentPayload): + raise TypeError( + f"{operation} requires MlaLatentPayload, got " + f"{type(payload).__name__}." + ) + for name, tensor, width in ( + ("latent_cache", payload.latent_cache, self.spec.kv_lora_rank), + ("rope_cache", payload.rope_cache, self.spec.rope_dim), + ): + if tensor.device != self.device: + raise ValueError( + f"{name} is on {tensor.device}, expected {self.device}." + ) + if tensor.dtype != self.spec.cache_dtype: + raise TypeError( + f"{name} must use {self.spec.cache_dtype}, got {tensor.dtype}." + ) + if tensor.ndim != 3 or tuple(tensor.shape[1:]) != (1, width): + raise ValueError( + f"{name} must have shape [slots, 1, {width}], got " + f"{tuple(tensor.shape)}." + ) + if payload.latent_cache.shape[0] != payload.rope_cache.shape[0]: + raise ValueError("MLA latent and RoPE caches must have equal slots.") + return payload + + def _get_prefill_plan( + self, + meta: AttentionViewMeta, + *, + cache_slot_count: int, + ) -> _MlaPrefillPlan: + cached = self._prefill_plan + validation_scope = get_context().attention_validation_scope + if cached is not None and cached.matches( + validation_scope, + meta, + cache_slot_count, + ): + return cached + + metadata = { + "active_slots": meta.active_slots, + "req_indices": meta.req_indices, + "context_lens": meta.context_lens, + } + for name, tensor in metadata.items(): + if tensor.device != self.device: + raise ValueError( + f"{name} is on {tensor.device}, expected {self.device}." + ) + if tensor.dtype != torch.int32: + raise TypeError( + f"{name} must use {torch.int32}, got {tensor.dtype}." + ) + if meta.context_lens.ndim != 1 or meta.context_lens.numel() == 0: + raise ValueError( + "MLA prefill context_lens must be a non-empty 1D tensor." + ) + batch_size = int(meta.context_lens.numel()) + if meta.active_slots.ndim != 2: + raise ValueError( + "MLA prefill active_slots must have shape " + "[rows, max_context_len]." + ) + if meta.req_indices.shape != (batch_size,): + raise ValueError( + f"MLA prefill req_indices must have shape ({batch_size},), " + f"got {tuple(meta.req_indices.shape)}." + ) + if batch_size > self.max_batch_size: + raise ValueError( + "MLA prefill batch exceeds the bound operator capacity: " + f"batch={batch_size} max_batch_size={self.max_batch_size}." + ) + + lengths = [int(value) for value in meta.context_lens.tolist()] + if any(length < 0 for length in lengths): + raise ValueError( + f"MLA prefill context lengths must be non-negative: {lengths}." + ) + total_visible_tokens = int(sum(lengths)) + if total_visible_tokens <= 0: + raise ValueError("MLA prefill requires at least one visible token.") + max_context_len = int(max(lengths)) + if ( + meta.max_context_len is not None + and int(meta.max_context_len) < max_context_len + ): + raise ValueError( + "MLA prefill max_context_len is smaller than an actual context: " + f"declared={meta.max_context_len} actual={max_context_len}." + ) + required_bytes = estimate_mla_prefill_workspace_bytes( + total_visible_tokens=total_visible_tokens, + batch_size=batch_size, + max_context_len=max_context_len, + local_q_heads=self.spec.local_q_heads, + activation_dtype=self.spec.activation_dtype, + cache_dtype=self.spec.cache_dtype, + ) + if required_bytes > self.prefill_workspace_bytes: + raise MemoryError( + "MLA full-history prefill workspace exceeds its configured " + f"budget: required={required_bytes} bytes budget=" + f"{self.prefill_workspace_bytes} bytes visible_tokens=" + f"{total_visible_tokens} local_heads={self.spec.local_q_heads}." + ) + + packed_starts: list[int] = [] + cursor = 0 + for length in lengths: + packed_starts.append(cursor) + cursor += length + packed_offsets = torch.tensor( + packed_starts, + dtype=torch.int32, + device=self.device, + ) + local_req_indices = torch.arange( + batch_size, + dtype=torch.int32, + device=self.device, + ) + positions = torch.arange( + max_context_len, + dtype=torch.int32, + device=self.device, + ) + packed_slots = packed_offsets[:, None] + positions[None, :] + validate_gather_metadata( + meta.active_slots, + meta.req_indices, + meta.context_lens, + packed_offsets, + cache_slot_count=cache_slot_count, + output_capacity=total_visible_tokens, + max_context_len=max_context_len, + ) + plan = _MlaPrefillPlan( + validation_scope=validation_scope, + source_active_slots=meta.active_slots, + source_req_indices=meta.req_indices, + source_context_lens=meta.context_lens, + source_max_context_len=meta.max_context_len, + cache_slot_count=int(cache_slot_count), + packed_offsets=packed_offsets, + packed_slots=packed_slots, + local_req_indices=local_req_indices, + total_visible_tokens=total_visible_tokens, + max_context_len=max_context_len, + required_workspace_bytes=required_bytes, + ) + self._prefill_plan = plan + return plan + + def prepare_prefill_history( + self, + view: PrefillComputeView, + ) -> MlaPrefillHistory: + if not isinstance(view, PrefillComputeView): + raise TypeError( + "MLA prefill requires PrefillComputeView, got " + f"{type(view).__name__}." + ) + payload = self._require_mla_payload(view, operation="MLA prefill") + meta = view.meta + plan = self._get_prefill_plan( + meta, + cache_slot_count=int(payload.latent_cache.shape[0]), + ) + gathered_latent = torch.empty( + plan.total_visible_tokens, + self.spec.kv_lora_rank, + dtype=self.spec.cache_dtype, + device=self.device, + ) + gathered_rope = torch.empty( + plan.total_visible_tokens, + self.spec.rope_dim, + dtype=self.spec.cache_dtype, + device=self.device, + ) + gather_latent_history( + payload.latent_cache, + payload.rope_cache, + meta.active_slots, + meta.req_indices, + meta.context_lens, + plan.packed_offsets, + gathered_latent, + gathered_rope, + max_context_len=plan.max_context_len, + validate_metadata=False, + ) + return MlaPrefillHistory( + gathered_latent=gathered_latent, + gathered_rope=gathered_rope, + packed_offsets=plan.packed_offsets, + packed_slots=plan.packed_slots, + local_req_indices=plan.local_req_indices, + context_lens=meta.context_lens, + max_context_len=plan.max_context_len, + required_workspace_bytes=plan.required_workspace_bytes, + ) + + def bind_prefill_kv( + self, + history: MlaPrefillHistory, + *, + expanded_k: torch.Tensor, + expanded_v: torch.Tensor, + ) -> MlaPrefillWorkset: + if not isinstance(history, MlaPrefillHistory): + raise TypeError( + "bind_prefill_kv requires MlaPrefillHistory, got " + f"{type(history).__name__}." + ) + if ( + history.gathered_latent.device != self.device + or history.gathered_rope.device != self.device + ): + raise ValueError("MLA prefill history is on the wrong device.") + if ( + history.gathered_latent.dtype != self.spec.cache_dtype + or history.gathered_rope.dtype != self.spec.cache_dtype + ): + raise TypeError("MLA prefill history uses the wrong cache dtype.") + expected_k_shape = ( + history.visible_tokens, + self.spec.local_q_heads, + self.spec.qk_head_dim, + ) + expected_v_shape = ( + history.visible_tokens, + self.spec.local_q_heads, + self.spec.value_head_dim, + ) + for name, tensor, expected_shape in ( + ("expanded_k", expanded_k, expected_k_shape), + ("expanded_v", expanded_v, expected_v_shape), + ): + if tuple(tensor.shape) != expected_shape: + raise ValueError( + f"{name} must have shape {expected_shape}, got " + f"{tuple(tensor.shape)}." + ) + if tensor.device != self.device: + raise ValueError( + f"{name} is on {tensor.device}, expected {self.device}." + ) + if tensor.dtype != self.spec.activation_dtype: + raise TypeError( + f"{name} must use {self.spec.activation_dtype}, got " + f"{tensor.dtype}." + ) + if tensor.stride(-1) != 1: + raise ValueError(f"{name} must be contiguous in its last dimension.") + return MlaPrefillWorkset( + history=history, + expanded_k=expanded_k, + expanded_v=expanded_v, + ) + + def run_prefill( + self, + q: torch.Tensor, + workset: MlaPrefillWorkset, + *, + b_start_loc: torch.Tensor, + chunk_lens: torch.Tensor, + ) -> torch.Tensor: + history = workset.history + if q.ndim != 3: + raise ValueError( + "MLA prefill q must have shape [tokens, local_heads, 256], " + f"got {tuple(q.shape)}." + ) + expected_q_shape = ( + int(q.shape[0]), + self.spec.local_q_heads, + self.spec.qk_head_dim, + ) + if tuple(q.shape) != expected_q_shape: + raise ValueError( + f"MLA prefill q must have shape {expected_q_shape}, got " + f"{tuple(q.shape)}." + ) + if q.device != self.device or q.dtype != self.spec.activation_dtype: + raise TypeError( + "MLA prefill q must match the operator device/dtype: " + f"q={q.device}/{q.dtype} expected=" + f"{self.device}/{self.spec.activation_dtype}." + ) + batch_size = int(history.context_lens.numel()) + for name, tensor in ( + ("b_start_loc", b_start_loc), + ("chunk_lens", chunk_lens), + ): + if tensor.shape != (batch_size,): + raise ValueError( + f"{name} must have shape ({batch_size},), got " + f"{tuple(tensor.shape)}." + ) + if tensor.device != self.device or tensor.dtype != torch.int32: + raise TypeError( + f"{name} must be int32 on {self.device}, got " + f"{tensor.device}/{tensor.dtype}." + ) + chunks = [int(value) for value in chunk_lens.tolist()] + starts = [int(value) for value in b_start_loc.tolist()] + expected_starts: list[int] = [] + cursor = 0 + for chunk in chunks: + expected_starts.append(cursor) + cursor += chunk + if starts != expected_starts or cursor != int(q.shape[0]): + raise ValueError( + "MLA prefill query packing is inconsistent: " + f"starts={starts} expected_starts={expected_starts} " + f"chunk_tokens={cursor} q_tokens={int(q.shape[0])}." + ) + contexts = [int(value) for value in history.context_lens.tolist()] + if any(chunk <= 0 or chunk > context for chunk, context in zip(chunks, contexts)): + raise ValueError( + "MLA prefill chunk lengths must be positive and no larger than " + f"their contexts: chunks={chunks} contexts={contexts}." + ) + + explicit_view = PrefillComputeView( + meta=AttentionViewMeta( + active_slots=history.packed_slots, + req_indices=history.local_req_indices, + context_lens=history.context_lens, + max_context_len=history.max_context_len, + ), + payload=ExplicitKVPayload( + k_cache=workset.expanded_k, + v_cache=workset.expanded_v, + ), + ) + return self.prefill_backend.run_prefill( + q, + explicit_view, + b_start_loc=b_start_loc, + chunk_lens=chunk_lens, + max_input_len=history.max_context_len, + ) + + def run_decode( + self, + q_nope_absorbed: torch.Tensor, + q_rope: torch.Tensor, + view: DecodeComputeView, + ) -> torch.Tensor: + if not isinstance(view, DecodeComputeView): + raise TypeError( + "MLA decode requires DecodeComputeView, got " + f"{type(view).__name__}." + ) + self._require_mla_payload(view, operation="MLA decode") + output = torch.empty_like(q_nope_absorbed) + return self.provider.run( + q_nope_absorbed, + q_rope, + view, + output, + validation_scope=get_context().attention_validation_scope, + ) + + +__all__ = [ + "MLAAttention", + "MlaPrefillHistory", + "MlaPrefillWorkset", + "estimate_mla_prefill_workspace_bytes", +] diff --git a/src/sparsevllm/operators/mla_attention.py b/src/sparsevllm/operators/mla_attention.py index 2f9a6806..84ca5003 100644 --- a/src/sparsevllm/operators/mla_attention.py +++ b/src/sparsevllm/operators/mla_attention.py @@ -20,6 +20,7 @@ MlaDecodeLaunchConfig, allocate_mla_decode_workspace, run_mla_decode, + validate_mla_decode_metadata, ) @@ -82,6 +83,8 @@ def run( q_rope: torch.Tensor, view: DecodeComputeView, output: torch.Tensor, + *, + validation_scope: object | None = None, ) -> torch.Tensor: raise NotImplementedError @@ -123,6 +126,13 @@ def __init__( config=self.launch_config, ) self.device = self.workspace.block_size.device + self._validated_decode_metadata: tuple[ + object, + torch.Tensor, + torch.Tensor, + torch.Tensor, + int, + ] | None = None @classmethod def supports( @@ -261,6 +271,8 @@ def run( q_rope: torch.Tensor, view: DecodeComputeView, output: torch.Tensor, + *, + validation_scope: object | None = None, ) -> torch.Tensor: payload = self._validate_run_inputs( q_nope_absorbed, @@ -268,6 +280,34 @@ def run( view, output, ) + cache_slot_count = int(payload.latent_cache.shape[0]) + metadata_key = ( + validation_scope, + view.meta.active_slots, + view.meta.req_indices, + view.meta.context_lens, + cache_slot_count, + ) + cached_key = self._validated_decode_metadata + metadata_is_validated = ( + validation_scope is not None + and cached_key is not None + and cached_key[0] is validation_scope + and cached_key[1] is metadata_key[1] + and cached_key[2] is metadata_key[2] + and cached_key[3] is metadata_key[3] + and cached_key[4] == metadata_key[4] + ) + if not metadata_is_validated: + validate_mla_decode_metadata( + view.meta.active_slots, + view.meta.req_indices, + view.meta.context_lens, + cache_slot_count=cache_slot_count, + ) + self._validated_decode_metadata = ( + metadata_key if validation_scope is not None else None + ) return run_mla_decode( q_nope_absorbed, q_rope, @@ -280,6 +320,7 @@ def run( self.workspace, softmax_scale=self.spec.softmax_scale, config=self.launch_config, + validate_metadata=False, ) diff --git a/src/sparsevllm/utils/context.py b/src/sparsevllm/utils/context.py index 5ea917a9..79678e3c 100644 --- a/src/sparsevllm/utils/context.py +++ b/src/sparsevllm/utils/context.py @@ -1,5 +1,6 @@ class Context: def __init__(self): + self.attention_validation_scope = object() self.is_prefill = False self.is_long_text = False self.cu_seqlens_q = None diff --git a/tests/test_attention_cache_storage.py b/tests/test_attention_cache_storage.py index b0bd2369..8babfec9 100644 --- a/tests/test_attention_cache_storage.py +++ b/tests/test_attention_cache_storage.py @@ -65,6 +65,12 @@ def test_glm_config_selects_mla_latent_layout(): config = _glm_config() assert config.attention_cache_layout == CacheLayout.MLA_LATENT.value + assert config.mla_prefill_workspace_bytes == 2 * 1024**3 + + +def test_mla_prefill_workspace_budget_must_be_positive(): + with pytest.raises(ValueError, match="mla_prefill_workspace_bytes"): + _glm_config(mla_prefill_workspace_bytes=0) @pytest.mark.parametrize( diff --git a/tests/test_mla_attention_layer.py b/tests/test_mla_attention_layer.py new file mode 100644 index 00000000..8b953814 --- /dev/null +++ b/tests/test_mla_attention_layer.py @@ -0,0 +1,367 @@ +from __future__ import annotations + +from unittest.mock import patch + +import pytest +import torch + +from sparsevllm.engine.cache_manager import ( + AttentionViewMeta, + ExplicitKVPayload, + MlaLatentPayload, + PrefillComputeView, +) +from sparsevllm.layers.mla_attention import ( + MLAAttention, + estimate_mla_prefill_workspace_bytes, +) +from sparsevllm.operators.mla_attention import ( + MlaAttentionOpSpec, + MlaAttentionProvider, +) +from sparsevllm.utils.context import reset_context + + +class _TestProvider(MlaAttentionProvider): + name = "test" + priority = 0 + + def __init__( + self, + spec: MlaAttentionOpSpec, + *, + device: torch.device | str, + max_batch_size: int, + ) -> None: + self.spec = spec + self.device = torch.device(device) + self.max_batch_size = int(max_batch_size) + + +def _spec(tp_size: int = 4) -> MlaAttentionOpSpec: + return MlaAttentionOpSpec( + num_q_heads=20, + kv_lora_rank=512, + rope_dim=64, + qk_head_dim=256, + value_head_dim=256, + activation_dtype=torch.bfloat16, + cache_dtype=torch.bfloat16, + tp_size=tp_size, + cuda_graph=False, + ) + + +def _attention( + *, + device: torch.device | str = "cpu", + tp_size: int = 4, + max_batch_size: int = 4, + budget: int = 64 * 1024 * 1024, +) -> MLAAttention: + spec = _spec(tp_size) + resolved_device = torch.device("cuda:0") if str(device) == "cuda" else device + return MLAAttention( + spec=spec, + provider=_TestProvider( + spec, + device=resolved_device, + max_batch_size=max_batch_size, + ), + prefill_workspace_bytes=budget, + ) + + +def _view( + latent_cache: torch.Tensor, + rope_cache: torch.Tensor, + active_slots: torch.Tensor, + request_indices: torch.Tensor, + context_lens: torch.Tensor, +) -> PrefillComputeView: + return PrefillComputeView( + meta=AttentionViewMeta( + active_slots=active_slots, + req_indices=request_indices, + context_lens=context_lens, + max_context_len=int(context_lens.max().item()), + ), + payload=MlaLatentPayload( + latent_cache=latent_cache, + rope_cache=rope_cache, + ), + ) + + +def _expand_history(attention: MLAAttention, history): + heads = attention.spec.local_q_heads + latent = history.gathered_latent + rope = history.gathered_rope[:, None, :].expand(-1, heads, -1) + k_nope = latent[:, None, :192].expand(-1, heads, -1) + expanded_k = torch.cat((k_nope, rope), dim=-1) + expanded_v = latent[:, None, 192:448].expand(-1, heads, -1).contiguous() + return attention.bind_prefill_kv( + history, + expanded_k=expanded_k, + expanded_v=expanded_v, + ) + + +def _torch_prefill( + q: torch.Tensor, + workset, + chunk_lens: torch.Tensor, +) -> torch.Tensor: + history = workset.history + outputs = [] + query_start = 0 + for batch_index, chunk_len in enumerate(chunk_lens.tolist()): + context_len = int(history.context_lens[batch_index].item()) + prefix_len = context_len - int(chunk_len) + history_start = int(history.packed_offsets[batch_index].item()) + keys = workset.expanded_k[ + history_start : history_start + context_len + ].float() + values = workset.expanded_v[ + history_start : history_start + context_len + ].float() + queries = q[query_start : query_start + chunk_len].float() + for query_offset, query in enumerate(queries): + visible = prefix_len + query_offset + 1 + logits = torch.einsum("hd,thd->ht", query, keys[:visible]) + probabilities = torch.softmax(logits * (256**-0.5), dim=-1) + outputs.append( + torch.einsum( + "ht,thd->hd", + probabilities, + values[:visible], + ).to(torch.bfloat16) + ) + query_start += int(chunk_len) + return torch.stack(outputs) + + +def test_mla_prefill_workspace_estimate_accounts_for_full_history() -> None: + actual = estimate_mla_prefill_workspace_bytes( + total_visible_tokens=11, + batch_size=2, + max_context_len=7, + local_q_heads=5, + activation_dtype=torch.bfloat16, + cache_dtype=torch.bfloat16, + ) + + gathered = 11 * (512 + 64) * 2 + expanded = 11 * 5 * ((192 + 256) + 256) * 2 + metadata = (2 * 7 + 2 * 2 + 7) * 4 + assert actual == gathered + expanded + metadata + + +def test_mla_prefill_rejects_wrong_payload_before_gather() -> None: + attention = _attention() + view = PrefillComputeView( + meta=AttentionViewMeta( + active_slots=torch.tensor([[0]], dtype=torch.int32), + req_indices=torch.tensor([0], dtype=torch.int32), + context_lens=torch.tensor([1], dtype=torch.int32), + ), + payload=ExplicitKVPayload( + k_cache=torch.empty(1, 1, 256, dtype=torch.bfloat16), + v_cache=torch.empty(1, 1, 256, dtype=torch.bfloat16), + ), + ) + + with ( + patch("sparsevllm.layers.mla_attention.gather_latent_history") as gather, + pytest.raises(TypeError, match="MlaLatentPayload"), + ): + attention.prepare_prefill_history(view) + gather.assert_not_called() + + +def test_mla_prefill_budget_fails_before_allocation_or_gather() -> None: + attention = _attention(budget=1) + view = _view( + torch.empty(2, 1, 512, dtype=torch.bfloat16), + torch.empty(2, 1, 64, dtype=torch.bfloat16), + torch.tensor([[0, 1]], dtype=torch.int32), + torch.tensor([0], dtype=torch.int32), + torch.tensor([2], dtype=torch.int32), + ) + + with ( + patch("sparsevllm.layers.mla_attention.gather_latent_history") as gather, + pytest.raises(MemoryError, match="exceeds its configured budget"), + ): + attention.prepare_prefill_history(view) + gather.assert_not_called() + + +def test_mla_attention_bind_resolves_provider_once() -> None: + spec = _spec() + provider = _TestProvider(spec, device="cpu", max_batch_size=8) + + with patch( + "sparsevllm.layers.mla_attention.resolve_mla_attention_provider", + return_value=provider, + ) as resolve: + attention = MLAAttention.bind( + spec=spec, + device="cpu", + max_batch_size=8, + prefill_workspace_bytes=1024, + ) + + assert attention.provider is provider + resolve.assert_called_once_with( + spec, + device="cpu", + max_batch_size=8, + ) + + +def test_mla_prefill_reuses_validated_packing_across_layers() -> None: + attention = _attention() + active_slots = torch.tensor([[3, 1]], dtype=torch.int32) + request_indices = torch.tensor([0], dtype=torch.int32) + context_lens = torch.tensor([2], dtype=torch.int32) + first_view = _view( + torch.empty(4, 1, 512, dtype=torch.bfloat16), + torch.empty(4, 1, 64, dtype=torch.bfloat16), + active_slots, + request_indices, + context_lens, + ) + second_view = _view( + torch.empty(4, 1, 512, dtype=torch.bfloat16), + torch.empty(4, 1, 64, dtype=torch.bfloat16), + active_slots, + request_indices, + context_lens, + ) + + with ( + patch( + "sparsevllm.layers.mla_attention.validate_gather_metadata" + ) as validate, + patch( + "sparsevllm.layers.mla_attention.gather_latent_history" + ) as gather, + ): + first = attention.prepare_prefill_history(first_view) + second = attention.prepare_prefill_history(second_view) + reset_context() + attention.prepare_prefill_history(second_view) + + assert validate.call_count == 2 + assert gather.call_count == 3 + assert first.packed_offsets is second.packed_offsets + assert first.packed_slots is second.packed_slots + + +CUDA_REQUIRED = pytest.mark.skipif( + not torch.cuda.is_available(), + reason="CUDA is required for MLA full-history prefill tests", +) + + +@CUDA_REQUIRED +def test_mla_prefill_matches_ragged_full_history_oracle() -> None: + torch.manual_seed(41) + attention = _attention(device="cuda", tp_size=4) + latent_cache = torch.randn(20, 1, 512, dtype=torch.bfloat16, device="cuda") + rope_cache = torch.randn(20, 1, 64, dtype=torch.bfloat16, device="cuda") + active_slots = torch.full((3, 7), -1, dtype=torch.int32, device="cuda") + active_slots[2, :5] = torch.tensor([13, 2, 17, 5, 11], device="cuda") + active_slots[0, :7] = torch.tensor( + [19, 1, 7, 15, 3, 9, 6], + device="cuda", + ) + context_lens = torch.tensor([5, 7], dtype=torch.int32, device="cuda") + view = _view( + latent_cache, + rope_cache, + active_slots, + torch.tensor([2, 0], dtype=torch.int32, device="cuda"), + context_lens, + ) + history = attention.prepare_prefill_history(view) + workset = _expand_history(attention, history) + chunk_lens = torch.tensor([2, 3], dtype=torch.int32, device="cuda") + b_start_loc = torch.tensor([0, 2], dtype=torch.int32, device="cuda") + q = torch.randn(5, 5, 256, dtype=torch.bfloat16, device="cuda") + + output = attention.run_prefill( + q, + workset, + b_start_loc=b_start_loc, + chunk_lens=chunk_lens, + ) + torch.cuda.synchronize() + expected = _torch_prefill(q, workset, chunk_lens) + + assert history.visible_tokens == 12 + torch.testing.assert_close( + output.float(), + expected.float(), + rtol=3e-2, + atol=3e-2, + ) + + +@CUDA_REQUIRED +def test_mla_prefill_is_invariant_to_chunk_boundary() -> None: + torch.manual_seed(43) + attention = _attention(device="cuda", tp_size=4) + latent_cache = torch.randn(12, 1, 512, dtype=torch.bfloat16, device="cuda") + rope_cache = torch.randn(12, 1, 64, dtype=torch.bfloat16, device="cuda") + active_slots = torch.tensor( + [[9, 1, 11, 3, 7, 5]], + dtype=torch.int32, + device="cuda", + ) + request_indices = torch.tensor([0], dtype=torch.int32, device="cuda") + q = torch.randn(6, 5, 256, dtype=torch.bfloat16, device="cuda") + + full_history = attention.prepare_prefill_history( + _view( + latent_cache, + rope_cache, + active_slots, + request_indices, + torch.tensor([6], dtype=torch.int32, device="cuda"), + ) + ) + full_workset = _expand_history(attention, full_history) + full_output = attention.run_prefill( + q, + full_workset, + b_start_loc=torch.tensor([0], dtype=torch.int32, device="cuda"), + chunk_lens=torch.tensor([6], dtype=torch.int32, device="cuda"), + ) + + first_history = attention.prepare_prefill_history( + _view( + latent_cache, + rope_cache, + active_slots, + request_indices, + torch.tensor([4], dtype=torch.int32, device="cuda"), + ) + ) + first_output = attention.run_prefill( + q[:4], + _expand_history(attention, first_history), + b_start_loc=torch.tensor([0], dtype=torch.int32, device="cuda"), + chunk_lens=torch.tensor([4], dtype=torch.int32, device="cuda"), + ) + second_output = attention.run_prefill( + q[4:], + full_workset, + b_start_loc=torch.tensor([0], dtype=torch.int32, device="cuda"), + chunk_lens=torch.tensor([2], dtype=torch.int32, device="cuda"), + ) + torch.cuda.synchronize() + + torch.testing.assert_close(first_output, full_output[:4], rtol=3e-2, atol=3e-2) + torch.testing.assert_close(second_output, full_output[4:], rtol=3e-2, atol=3e-2) diff --git a/tests/test_mla_attention_operator.py b/tests/test_mla_attention_operator.py index db39bdd3..3bf28897 100644 --- a/tests/test_mla_attention_operator.py +++ b/tests/test_mla_attention_operator.py @@ -206,6 +206,7 @@ def test_mla_provider_run_does_not_resolve_or_allocate() -> None: q_nope_absorbed = torch.empty(2, 5, 512, dtype=torch.bfloat16) q_rope = torch.empty(2, 5, 64, dtype=torch.bfloat16) output = torch.empty_like(q_nope_absorbed) + validation_scope = object() with ( patch.object(OpResolver, "resolve") as resolve, @@ -216,13 +217,44 @@ def test_mla_provider_run_does_not_resolve_or_allocate() -> None: "sparsevllm.operators.mla_attention.run_mla_decode", return_value=output, ) as kernel, + patch( + "sparsevllm.operators.mla_attention.validate_mla_decode_metadata" + ) as validate, ): - actual = provider.run(q_nope_absorbed, q_rope, view, output) + actual = provider.run( + q_nope_absorbed, + q_rope, + view, + output, + validation_scope=validation_scope, + ) + provider.run( + q_nope_absorbed, + q_rope, + view, + output, + validation_scope=validation_scope, + ) + provider.run( + q_nope_absorbed, + q_rope, + view, + output, + validation_scope=object(), + ) assert actual is output resolve.assert_not_called() allocate.assert_not_called() - kernel.assert_called_once_with( + assert validate.call_count == 2 + validate.assert_called_with( + view.meta.active_slots, + view.meta.req_indices, + view.meta.context_lens, + cache_slot_count=4, + ) + assert kernel.call_count == 3 + kernel.assert_called_with( q_nope_absorbed, q_rope, payload.latent_cache, @@ -234,6 +266,7 @@ def test_mla_provider_run_does_not_resolve_or_allocate() -> None: workspace, softmax_scale=spec.softmax_scale, config=provider.launch_config, + validate_metadata=False, ) From 4551cd3e73059c5806029695cfd733e2c34fc3a0 Mon Sep 17 00:00:00 2001 From: Codex Date: Tue, 4 Aug 2026 23:00:17 +0800 Subject: [PATCH 06/48] refactor: share packed moe experts --- src/sparsevllm/layers/packed_moe.py | 305 ++++++++++++++++++++++++++++ src/sparsevllm/models/qwen3_moe.py | 272 ++----------------------- tests/test_packed_moe.py | 69 +++++++ 3 files changed, 390 insertions(+), 256 deletions(-) create mode 100644 src/sparsevllm/layers/packed_moe.py create mode 100644 tests/test_packed_moe.py diff --git a/src/sparsevllm/layers/packed_moe.py b/src/sparsevllm/layers/packed_moe.py new file mode 100644 index 00000000..0356c17d --- /dev/null +++ b/src/sparsevllm/layers/packed_moe.py @@ -0,0 +1,305 @@ +from __future__ import annotations + +from collections.abc import Callable + +import torch +from torch import nn + +from sparsevllm.distributed import get_parallel_context +from sparsevllm.layers.expert_weights import ( + PackedExpertWeightLoader, + UnquantizedExpertTpShard, +) +from sparsevllm.operators.moe import ( + MoeOpSpec, + MoeProvider, + resolve_moe_provider, +) +from sparsevllm.quantization.fp8_tp import Fp8ExpertTpShard + + +class PackedMoeExperts(PackedExpertWeightLoader, nn.Module): + """Shared packed expert storage, execution, and checkpoint loading.""" + + checkpoint_projection_map = { + "gate_proj": "gate", + "up_proj": "up", + "down_proj": "down", + } + + def __init__( + self, + *, + num_experts: int, + hidden_size: int, + intermediate_size: int, + top_k: int, + activation_dtype: torch.dtype, + fp8_enabled: bool, + cuda_graph: bool, + routing_method: str = "softmax", + scale_dtype: torch.dtype | None = None, + model_label: str = "PackedMoE", + provider_resolver: Callable[[MoeOpSpec], MoeProvider] = resolve_moe_provider, + parallel_context=None, + ) -> None: + super().__init__() + self.model_label = str(model_label) + if parallel_context is None: + parallel_context = get_parallel_context() + self.tp_rank = parallel_context.moe_tp_rank + self.tp_size = parallel_context.moe_tp_size + self.ep_rank = parallel_context.ep_rank + self.ep_size = parallel_context.ep_size + self.num_experts = int(num_experts) + self.hidden_size = int(hidden_size) + self.global_intermediate_size = int(intermediate_size) + if self.global_intermediate_size % self.tp_size: + raise ValueError( + f"{self.model_label} intermediate_size must be divisible by " + f"MoE tensor parallel size, got " + f"{self.global_intermediate_size} and {self.tp_size}." + ) + self.logical_intermediate_size = ( + self.global_intermediate_size // self.tp_size + ) + self.fp8_enabled = bool(fp8_enabled) + if self.fp8_enabled and ( + self.hidden_size % 128 or self.global_intermediate_size % 128 + ): + raise ValueError( + f"{self.model_label} FP8 requires hidden/intermediate sizes " + "aligned to 128, got " + f"{self.hidden_size}/{self.global_intermediate_size}." + ) + self.fp8_tp_shard = ( + Fp8ExpertTpShard( + self.global_intermediate_size, + self.tp_rank, + self.tp_size, + ) + if self.fp8_enabled + else None + ) + self.checkpoint_tp_shard = ( + self.fp8_tp_shard + if self.fp8_tp_shard is not None + else UnquantizedExpertTpShard( + self.global_intermediate_size, + self.tp_rank, + self.tp_size, + ) + ) + self.intermediate_size = ( + self.fp8_tp_shard.physical_size + if self.fp8_tp_shard is not None + else self.logical_intermediate_size + ) + if self.num_experts <= 0 or self.num_experts % self.ep_size: + raise ValueError( + f"{self.model_label} num_experts must be positive and divisible " + f"by EP size, got {self.num_experts} and {self.ep_size}." + ) + self.num_local_experts = self.num_experts // self.ep_size + self.local_expert_start = self.ep_rank * self.num_local_experts + self.local_expert_end = self.local_expert_start + self.num_local_experts + self.op_spec = MoeOpSpec( + num_experts=self.num_experts, + num_local_experts=self.num_local_experts, + hidden_size=self.hidden_size, + intermediate_size=self.intermediate_size, + top_k=int(top_k), + activation_dtype=activation_dtype, + weight_dtype=( + torch.float8_e4m3fn if self.fp8_enabled else activation_dtype + ), + block_shape=(128, 128) if self.fp8_enabled else None, + ep_size=int(self.ep_size), + cuda_graph=bool(cuda_graph), + tp_size=int(self.tp_size), + routing_method=str(routing_method), + scale_dtype=scale_dtype, + ) + self.provider = provider_resolver(self.op_spec) + self.w13_weight = nn.Parameter( + torch.empty( + self.num_local_experts, + 2 * self.intermediate_size, + self.hidden_size, + dtype=torch.float8_e4m3fn if self.fp8_enabled else None, + ), + requires_grad=not self.fp8_enabled, + ) + self.w2_weight = nn.Parameter( + torch.empty( + self.num_local_experts, + self.hidden_size, + self.intermediate_size, + dtype=torch.float8_e4m3fn if self.fp8_enabled else None, + ), + requires_grad=not self.fp8_enabled, + ) + if self.fp8_enabled: + self.register_buffer( + "w13_scale_inv", + torch.empty( + self.num_local_experts, + 2 * self.intermediate_size // 128, + self.hidden_size // 128, + dtype=torch.float32, + ), + ) + self.register_buffer( + "w2_scale_inv", + torch.empty( + self.num_local_experts, + self.hidden_size // 128, + self.intermediate_size // 128, + dtype=torch.float32, + ), + ) + else: + self.register_buffer("w13_scale_inv", None) + self.register_buffer("w2_scale_inv", None) + self._loaded_expert_shards: set[tuple[int, str]] = set() + + def is_local_expert(self, global_expert_id: int) -> bool: + return self.local_expert_start <= int(global_expert_id) < self.local_expert_end + + def load_expert_weight( + self, + global_expert_id: int, + projection: str, + loaded_weight: torch.Tensor, + loaded_scale: torch.Tensor | None = None, + ) -> None: + global_expert_id = int(global_expert_id) + if not self.is_local_expert(global_expert_id): + raise ValueError( + f"Expert {global_expert_id} is outside local range " + f"[{self.local_expert_start}, {self.local_expert_end})." + ) + logical_projection = self.checkpoint_projection_map.get(projection) + if logical_projection is None: + raise ValueError(f"Unsupported expert projection {projection!r}.") + load_key = (global_expert_id, projection) + if load_key in self._loaded_expert_shards: + raise ValueError( + f"Duplicate {self.model_label} expert weight for " + f"expert={global_expert_id}, projection={projection}." + ) + + if self.fp8_enabled: + if loaded_scale is None: + raise ValueError( + f"Missing FP8 weight_scale_inv for {self.model_label} " + f"expert={global_expert_id}, projection={projection}." + ) + if loaded_weight.dtype != torch.float8_e4m3fn: + raise TypeError( + f"{self.model_label} expert weight must be FP8 E4M3, " + f"got {loaded_weight.dtype}." + ) + if loaded_scale.dtype != torch.bfloat16: + raise TypeError( + f"{self.model_label} expert weight_scale_inv must be BF16, " + f"got {loaded_scale.dtype}." + ) + elif loaded_scale is not None: + raise ValueError( + f"Unexpected weight_scale_inv for unquantized {self.model_label} " + f"expert={global_expert_id}, projection={projection}." + ) + + if self.fp8_tp_shard is not None: + loaded_weight, loaded_scale = self.fp8_tp_shard.prepare_projection( + loaded_weight, + loaded_scale, + hidden_size=self.hidden_size, + down_projection=logical_projection == "down", + ) + else: + loaded_weight = self._local_projection_shard( + logical_projection, + loaded_weight, + ) + local_expert_id = global_expert_id - self.local_expert_start + self.provider.load_expert_projection( + self.op_spec, + local_expert_id=local_expert_id, + projection=logical_projection, + loaded_weight=loaded_weight, + loaded_scale=loaded_scale, + w13_weight=self.w13_weight.data, + w2_weight=self.w2_weight.data, + w13_scale_inv=self.w13_scale_inv, + w2_scale_inv=self.w2_scale_inv, + ) + self._loaded_expert_shards.add(load_key) + + def _local_projection_shard( + self, + logical_projection: str, + loaded_weight: torch.Tensor, + ) -> torch.Tensor: + down_projection = logical_projection == "down" + local_shape = ( + (self.hidden_size, self.intermediate_size) + if down_projection + else (self.intermediate_size, self.hidden_size) + ) + if tuple(loaded_weight.shape) == local_shape: + return loaded_weight + + global_shape = ( + (self.hidden_size, self.global_intermediate_size) + if down_projection + else (self.global_intermediate_size, self.hidden_size) + ) + if tuple(loaded_weight.shape) != global_shape: + raise ValueError( + f"{self.model_label} expert projection shape mismatch: " + f"projection={logical_projection}, expected local={local_shape} " + f"or global={global_shape}, got={tuple(loaded_weight.shape)}." + ) + shard_dim = 1 if down_projection else 0 + return loaded_weight.chunk(self.tp_size, dim=shard_dim)[self.tp_rank] + + def validate_loaded_weights(self) -> None: + expected = { + (global_expert_id, projection) + for global_expert_id in range( + self.local_expert_start, + self.local_expert_end, + ) + for projection in self.checkpoint_projection_map + } + missing = sorted(expected - self._loaded_expert_shards) + if missing: + raise ValueError( + f"Missing local {self.model_label} expert weights: " + f"local_range=[{self.local_expert_start}, " + f"{self.local_expert_end}), missing={missing[:8]}." + ) + + def forward( + self, + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + ) -> torch.Tensor: + return self.provider.run( + self.op_spec, + hidden_states, + topk_ids, + topk_weights, + self.w13_weight, + self.w2_weight, + self.w13_scale_inv, + self.w2_scale_inv, + local_expert_start=self.local_expert_start, + ep_rank=int(self.ep_rank), + ) + + +__all__ = ["PackedMoeExperts"] diff --git a/src/sparsevllm/models/qwen3_moe.py b/src/sparsevllm/models/qwen3_moe.py index 3a9d04d5..eefe0153 100644 --- a/src/sparsevllm/models/qwen3_moe.py +++ b/src/sparsevllm/models/qwen3_moe.py @@ -10,18 +10,13 @@ from sparsevllm.distributed import get_parallel_context from sparsevllm.layers.embed_head import ParallelLMHead -from sparsevllm.layers.expert_weights import ( - PackedExpertWeightLoader, - UnquantizedExpertTpShard, -) +from sparsevllm.layers.packed_moe import PackedMoeExperts from sparsevllm.models.qwen3 import Qwen3DecoderLayerBase, Qwen3ModelBase from sparsevllm.operators.moe import ( - MoeOpSpec, model_activation_dtype, resolve_moe_provider, ) from sparsevllm.platforms import device_runtime -from sparsevllm.quantization.fp8_tp import Fp8ExpertTpShard from sparsevllm.utils.log import logger from sparsevllm.utils.weight_target import WeightTarget @@ -61,260 +56,25 @@ def forward( return router_logits, topk_weights, topk_ids -class Qwen3MoePackedExperts(PackedExpertWeightLoader, nn.Module): - checkpoint_projection_map = { - "gate_proj": "gate", - "up_proj": "up", - "down_proj": "down", - } - +class Qwen3MoePackedExperts(PackedMoeExperts): def __init__(self, config: Qwen3MoeConfig) -> None: - super().__init__() - parallel_context = get_parallel_context() - self.tp_rank = parallel_context.moe_tp_rank - self.tp_size = parallel_context.moe_tp_size - self.ep_rank = parallel_context.ep_rank - self.ep_size = parallel_context.ep_size - self.num_experts = int(config.num_experts) - self.hidden_size = int(config.hidden_size) - self.global_intermediate_size = int(config.moe_intermediate_size) - if self.global_intermediate_size % self.tp_size: - raise ValueError( - "Qwen3MoE moe_intermediate_size must be divisible by " - f"tensor_parallel_size, got {self.global_intermediate_size} and " - f"{self.tp_size}." - ) - self.logical_intermediate_size = self.global_intermediate_size // self.tp_size - self.fp8_enabled = bool( - getattr(getattr(config, "quantization_config", None), "enabled", False) - ) - if self.fp8_enabled and ( - self.hidden_size % 128 or self.global_intermediate_size % 128 - ): - raise ValueError( - "Qwen3MoE FP8 requires hidden_size and moe_intermediate_size " - f"aligned to 128, got {self.hidden_size}/{self.global_intermediate_size}." - ) - self.fp8_tp_shard = ( - Fp8ExpertTpShard( - self.global_intermediate_size, - self.tp_rank, - self.tp_size, - ) - if self.fp8_enabled - else None - ) - self.checkpoint_tp_shard = ( - self.fp8_tp_shard - if self.fp8_tp_shard is not None - else UnquantizedExpertTpShard( - self.global_intermediate_size, - self.tp_rank, - self.tp_size, - ) - ) - self.intermediate_size = ( - self.fp8_tp_shard.physical_size - if self.fp8_tp_shard is not None - else self.logical_intermediate_size - ) - self.num_local_experts = self.num_experts // self.ep_size - self.local_expert_start = self.ep_rank * self.num_local_experts - self.local_expert_end = self.local_expert_start + self.num_local_experts - activation_dtype = model_activation_dtype(config) - self.op_spec = MoeOpSpec( - num_experts=self.num_experts, - num_local_experts=self.num_local_experts, - hidden_size=self.hidden_size, - intermediate_size=self.intermediate_size, + super().__init__( + num_experts=int(config.num_experts), + hidden_size=int(config.hidden_size), + intermediate_size=int(config.moe_intermediate_size), top_k=int(config.num_experts_per_tok), - activation_dtype=activation_dtype, - weight_dtype=( - torch.float8_e4m3fn if self.fp8_enabled else activation_dtype + activation_dtype=model_activation_dtype(config), + fp8_enabled=bool( + getattr( + getattr(config, "quantization_config", None), + "enabled", + False, + ) ), - block_shape=(128, 128) if self.fp8_enabled else None, - ep_size=int(self.ep_size), cuda_graph=bool(getattr(config, "decode_cuda_graph", False)), - tp_size=int(self.tp_size), - ) - self.provider = resolve_moe_provider(self.op_spec) - self.w13_weight = nn.Parameter( - torch.empty( - self.num_local_experts, - 2 * self.intermediate_size, - self.hidden_size, - dtype=torch.float8_e4m3fn if self.fp8_enabled else None, - ), - requires_grad=not self.fp8_enabled, - ) - self.w2_weight = nn.Parameter( - torch.empty( - self.num_local_experts, - self.hidden_size, - self.intermediate_size, - dtype=torch.float8_e4m3fn if self.fp8_enabled else None, - ), - requires_grad=not self.fp8_enabled, - ) - if self.fp8_enabled: - self.register_buffer( - "w13_scale_inv", - torch.empty( - self.num_local_experts, - 2 * self.intermediate_size // 128, - self.hidden_size // 128, - dtype=torch.float32, - ), - ) - self.register_buffer( - "w2_scale_inv", - torch.empty( - self.num_local_experts, - self.hidden_size // 128, - self.intermediate_size // 128, - dtype=torch.float32, - ), - ) - else: - self.register_buffer("w13_scale_inv", None) - self.register_buffer("w2_scale_inv", None) - self._loaded_expert_shards: set[tuple[int, str]] = set() - - def is_local_expert(self, global_expert_id: int) -> bool: - return self.local_expert_start <= int(global_expert_id) < self.local_expert_end - - def load_expert_weight( - self, - global_expert_id: int, - projection: str, - loaded_weight: torch.Tensor, - loaded_scale: torch.Tensor | None = None, - ) -> None: - global_expert_id = int(global_expert_id) - if not self.is_local_expert(global_expert_id): - raise ValueError( - f"Expert {global_expert_id} is outside local range " - f"[{self.local_expert_start}, {self.local_expert_end})." - ) - if projection not in {"gate_proj", "up_proj", "down_proj"}: - raise ValueError(f"Unsupported expert projection {projection!r}.") - load_key = (global_expert_id, projection) - if load_key in self._loaded_expert_shards: - raise ValueError( - f"Duplicate Qwen3MoE expert weight for expert={global_expert_id}, " - f"projection={projection}." - ) - - if self.fp8_enabled: - if loaded_scale is None: - raise ValueError( - "Missing FP8 weight_scale_inv for Qwen3MoE " - f"expert={global_expert_id}, projection={projection}." - ) - if loaded_weight.dtype != torch.float8_e4m3fn: - raise TypeError( - "Qwen3MoE expert weight must be FP8 E4M3, " - f"got {loaded_weight.dtype}." - ) - if loaded_scale.dtype != torch.bfloat16: - raise TypeError( - "Qwen3MoE expert weight_scale_inv must be BF16, " - f"got {loaded_scale.dtype}." - ) - elif loaded_scale is not None: - raise ValueError( - "Unexpected weight_scale_inv for unquantized Qwen3MoE " - f"expert={global_expert_id}, projection={projection}." - ) - - if self.fp8_tp_shard is not None: - loaded_weight, loaded_scale = self.fp8_tp_shard.prepare_projection( - loaded_weight, - loaded_scale, - hidden_size=self.hidden_size, - down_projection=projection == "down_proj", - ) - else: - loaded_weight = self._local_projection_shard(projection, loaded_weight) - local_expert_id = global_expert_id - self.local_expert_start - logical_projection = { - "gate_proj": "gate", - "up_proj": "up", - "down_proj": "down", - }[projection] - self.provider.load_expert_projection( - self.op_spec, - local_expert_id=local_expert_id, - projection=logical_projection, - loaded_weight=loaded_weight, - loaded_scale=loaded_scale, - w13_weight=self.w13_weight.data, - w2_weight=self.w2_weight.data, - w13_scale_inv=self.w13_scale_inv, - w2_scale_inv=self.w2_scale_inv, - ) - self._loaded_expert_shards.add(load_key) - - def _local_projection_shard( - self, - projection: str, - loaded_weight: torch.Tensor, - ) -> torch.Tensor: - local_shape = ( - (self.hidden_size, self.intermediate_size) - if projection == "down_proj" - else (self.intermediate_size, self.hidden_size) - ) - if tuple(loaded_weight.shape) == local_shape: - return loaded_weight - - global_shape = ( - (self.hidden_size, self.global_intermediate_size) - if projection == "down_proj" - else (self.global_intermediate_size, self.hidden_size) - ) - if tuple(loaded_weight.shape) != global_shape: - raise ValueError( - "Qwen3MoE expert projection shape mismatch: " - f"projection={projection}, expected local={local_shape} or " - f"global={global_shape}, got={tuple(loaded_weight.shape)}." - ) - shard_dim = 1 if projection == "down_proj" else 0 - return loaded_weight.chunk(self.tp_size, dim=shard_dim)[self.tp_rank] - - def validate_loaded_weights(self) -> None: - expected = { - (global_expert_id, projection) - for global_expert_id in range( - self.local_expert_start, self.local_expert_end - ) - for projection in ("gate_proj", "up_proj", "down_proj") - } - missing = sorted(expected - self._loaded_expert_shards) - if missing: - raise ValueError( - "Missing local Qwen3MoE expert weights: " - f"local_range=[{self.local_expert_start}, {self.local_expert_end}), " - f"missing={missing[:8]}." - ) - - def forward( - self, - hidden_states: torch.Tensor, - topk_ids: torch.Tensor, - topk_weights: torch.Tensor, - ) -> torch.Tensor: - return self.provider.run( - self.op_spec, - hidden_states, - topk_ids, - topk_weights, - self.w13_weight, - self.w2_weight, - self.w13_scale_inv, - self.w2_scale_inv, - local_expert_start=self.local_expert_start, - ep_rank=int(self.ep_rank), + model_label="Qwen3MoE", + provider_resolver=resolve_moe_provider, + parallel_context=get_parallel_context(), ) diff --git a/tests/test_packed_moe.py b/tests/test_packed_moe.py new file mode 100644 index 00000000..1d3247df --- /dev/null +++ b/tests/test_packed_moe.py @@ -0,0 +1,69 @@ +from types import SimpleNamespace + +import torch + +from sparsevllm.layers.packed_moe import PackedMoeExperts +from sparsevllm.models.qwen3_moe import Qwen3MoePackedExperts +from sparsevllm.operators.moe import TritonMoeProvider + + +def _parallel_context(*, tp_rank=0, tp_size=1, ep_rank=0, ep_size=1): + return SimpleNamespace( + moe_tp_rank=tp_rank, + moe_tp_size=tp_size, + ep_rank=ep_rank, + ep_size=ep_size, + ) + + +def _experts(**overrides) -> PackedMoeExperts: + values = { + "num_experts": 4, + "hidden_size": 8, + "intermediate_size": 8, + "top_k": 2, + "activation_dtype": torch.bfloat16, + "fp8_enabled": False, + "cuda_graph": False, + "routing_method": "biased_sigmoid", + "model_label": "TestGLM", + "provider_resolver": lambda spec: TritonMoeProvider(), + "parallel_context": _parallel_context(), + } + values.update(overrides) + return PackedMoeExperts(**values) + + +def test_qwen_packed_experts_use_shared_physical_module() -> None: + assert issubclass(Qwen3MoePackedExperts, PackedMoeExperts) + + +def test_packed_experts_accept_glm_router_contract_without_owning_router() -> None: + experts = _experts(num_experts=64, top_k=4) + + assert experts.op_spec.routing_method == "biased_sigmoid" + assert experts.op_spec.top_k == 4 + assert experts.w13_weight.shape == (64, 16, 8) + assert experts.w2_weight.shape == (64, 8, 8) + + +def test_packed_experts_load_and_validate_rank_local_projections() -> None: + experts = _experts( + num_experts=4, + parallel_context=_parallel_context(ep_rank=1, ep_size=2), + ) + for global_expert_id in range(2, 4): + for projection, shape in ( + ("gate_proj", (8, 8)), + ("up_proj", (8, 8)), + ("down_proj", (8, 8)), + ): + experts.load_expert_weight( + global_expert_id, + projection, + torch.full(shape, float(global_expert_id), dtype=torch.bfloat16), + ) + + experts.validate_loaded_weights() + assert experts.local_expert_start == 2 + assert experts.local_expert_end == 4 From 552a9d3901e2927c3becdffc4a39506699cee44d Mon Sep 17 00:00:00 2001 From: Codex Date: Tue, 4 Aug 2026 23:42:50 +0800 Subject: [PATCH 07/48] feat: add glm4 moe lite model --- configs/debug/glm4_moe_lite_tiny_random.json | 6 + src/sparsevllm/debug/tiny_random.py | 20 +- src/sparsevllm/engine/model_runner.py | 18 + src/sparsevllm/layers/rotary_embedding.py | 36 +- src/sparsevllm/method_registry.py | 7 + src/sparsevllm/models/glm4_moe_lite.py | 916 ++++++++++++++++++ src/sparsevllm/models/spec.py | 8 + .../triton_kernel/minimax_m2_router.py | 56 +- .../triton_kernel/moe_biased_sigmoid.py | 133 +++ tests/test_attention_cache_storage.py | 23 +- tests/test_glm4_moe_lite.py | 322 ++++++ tests/test_minimax_m2_router.py | 29 + 12 files changed, 1517 insertions(+), 57 deletions(-) create mode 100644 configs/debug/glm4_moe_lite_tiny_random.json create mode 100644 src/sparsevllm/models/glm4_moe_lite.py create mode 100644 src/sparsevllm/triton_kernel/moe_biased_sigmoid.py create mode 100644 tests/test_glm4_moe_lite.py diff --git a/configs/debug/glm4_moe_lite_tiny_random.json b/configs/debug/glm4_moe_lite_tiny_random.json new file mode 100644 index 00000000..aa71dbe3 --- /dev/null +++ b/configs/debug/glm4_moe_lite_tiny_random.json @@ -0,0 +1,6 @@ +{ + "hidden_size": 64, + "intermediate_size": 128, + "max_position_embeddings": 512, + "num_hidden_layers": 2 +} diff --git a/src/sparsevllm/debug/tiny_random.py b/src/sparsevllm/debug/tiny_random.py index 8a9fba23..c07ee770 100644 --- a/src/sparsevllm/debug/tiny_random.py +++ b/src/sparsevllm/debug/tiny_random.py @@ -116,11 +116,22 @@ def apply_tiny_random_overrides(hf_config: Any, path: str) -> dict[str, int]: ) hf_config.layer_types = list(layer_types[:num_layers]) + mlp_layer_types = getattr(hf_config, "mlp_layer_types", None) + if mlp_layer_types is not None: + num_layers = int(hf_config.num_hidden_layers) + if len(mlp_layer_types) < num_layers: + raise ValueError( + "Tiny random config cannot expand mlp_layer_types: " + f"requested={num_layers}, available={len(mlp_layer_types)}." + ) + hf_config.mlp_layer_types = list(mlp_layer_types[:num_layers]) + hidden_size = int(hf_config.hidden_size) num_heads = int(hf_config.num_attention_heads) num_kv_heads = int(hf_config.num_key_value_heads) head_dim = int(getattr(hf_config, "head_dim", hidden_size // num_heads)) - if hidden_size != num_heads * head_dim: + is_mla = str(getattr(hf_config, "model_type", "") or "") == "glm4_moe_lite" + if not is_mla and hidden_size != num_heads * head_dim: raise ValueError( "Tiny random config requires hidden_size == num_attention_heads * head_dim, " f"got {hidden_size} != {num_heads} * {head_dim}." @@ -179,7 +190,12 @@ def initialize_sparse_model( loaded_count = 0 loaded_parameter_names: set[str] = set() try: - for source_weight_name, loaded_weight in reference.state_dict().items(): + reference_state = reference.state_dict() + reference_items = reference_state.items() + reference_adapter = getattr(model, "iter_tiny_reference_weights", None) + if callable(reference_adapter): + reference_items = reference_adapter(reference_state) + for source_weight_name, loaded_weight in reference_items: param_name = _target_weight_name_for_model(model, source_weight_name) if param_name is None: continue diff --git a/src/sparsevllm/engine/model_runner.py b/src/sparsevllm/engine/model_runner.py index 7caf2f42..0d486071 100644 --- a/src/sparsevllm/engine/model_runner.py +++ b/src/sparsevllm/engine/model_runner.py @@ -44,6 +44,11 @@ except ImportError: Qwen3MoeForCausalLM = None +try: + from sparsevllm.models.glm4_moe_lite import Glm4MoeLiteForCausalLM +except ImportError: + Glm4MoeLiteForCausalLM = None + try: from sparsevllm.models.minimax_m2 import MiniMaxM2ForCausalLM except ImportError: @@ -160,6 +165,19 @@ def __init__( "decode_cuda_graph", bool(getattr(config, "decode_cuda_graph", False)), ) + if config.model_spec.runtime_class_name == "Glm4MoeLiteForCausalLM": + setattr(hf_config, "sparsevllm_device", self.device) + setattr( + hf_config, + "sparsevllm_max_batch_size", + max(config.max_num_seqs_in_batch, config.max_decoding_seqs), + ) + setattr( + hf_config, + "sparsevllm_mla_prefill_workspace_bytes", + config.mla_prefill_workspace_bytes, + ) + setattr(hf_config, "sparsevllm_expect_mtp_weights", not config.tiny_random) self.model = _create_model(hf_config, config.model_spec) if config.tiny_random: diff --git a/src/sparsevllm/layers/rotary_embedding.py b/src/sparsevllm/layers/rotary_embedding.py index 73c9e035..706943da 100644 --- a/src/sparsevllm/layers/rotary_embedding.py +++ b/src/sparsevllm/layers/rotary_embedding.py @@ -16,6 +16,20 @@ def apply_rotary_emb( return torch.cat((y1, y2), dim=-1).to(x.dtype) +def apply_interleaved_rotary_emb( + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, +) -> torch.Tensor: + """Rotate adjacent pairs and return the split-half layout used by GLM.""" + + x_even = x.float()[..., 0::2] + x_odd = x.float()[..., 1::2] + y_even = x_even * cos - x_odd * sin + y_odd = x_odd * cos + x_even * sin + return torch.cat((y_even, y_odd), dim=-1).to(x.dtype) + + def apply_partial_rotary_emb( rotary_emb: "RotaryEmbedding", positions: torch.Tensor, @@ -83,6 +97,7 @@ def __init__( base: float, rope_scaling: tuple[tuple[str, object], ...] | None = None, backend: str = "flashinfer", + interleaved: bool = False, ) -> None: super().__init__() if backend not in {"flashinfer", "torch"}: @@ -90,6 +105,7 @@ def __init__( f"Unsupported RoPE backend={backend!r}; expected 'flashinfer' or 'torch'." ) self.backend = backend + self.interleaved = bool(interleaved) self.head_size = head_size assert rotary_dim == head_size inv_freq = 1.0 / (base**(torch.arange(0, rotary_dim, 2, dtype=torch.float) / rotary_dim)) @@ -141,8 +157,13 @@ def compiled_forward( ) -> tuple[torch.Tensor, torch.Tensor]: cos_sin = self.cos_sin_cache[positions] cos, sin = cos_sin.chunk(2, dim=-1) - query = apply_rotary_emb(query, cos, sin) - key = apply_rotary_emb(key, cos, sin) + apply = ( + apply_interleaved_rotary_emb + if self.interleaved + else apply_rotary_emb + ) + query = apply(query, cos, sin) + key = apply(key, cos, sin) return query, key def flashinfer_forward( @@ -162,9 +183,14 @@ def flashinfer_forward( key=key_flat, head_size=head_size, cos_sin_cache=cos_sin_cache, - is_neox=True, + is_neox=not self.interleaved, ) - return query_out.view_as(query), key_out.view_as(key) + query_out = query_out.view_as(query) + key_out = key_out.view_as(key) + if self.interleaved: + query_out = query_out.unflatten(-1, (-1, 2)).transpose(-1, -2).flatten(-2) + key_out = key_out.unflatten(-1, (-1, 2)).transpose(-1, -2).flatten(-2) + return query_out, key_out def forward( self, @@ -185,6 +211,7 @@ def get_rope( base: float, rope_scaling: tuple[tuple[str, object], ...] | None = None, backend: str = "flashinfer", + interleaved: bool = False, ): rotary_emb = RotaryEmbedding( head_size, @@ -193,5 +220,6 @@ def get_rope( base, rope_scaling=rope_scaling, backend=backend, + interleaved=interleaved, ) return rotary_emb diff --git a/src/sparsevllm/method_registry.py b/src/sparsevllm/method_registry.py index ab280cac..98c0d693 100644 --- a/src/sparsevllm/method_registry.py +++ b/src/sparsevllm/method_registry.py @@ -89,6 +89,12 @@ class ModelRuntimeCompatibility: decode_cuda_graph_methods=frozenset(CANONICAL_SPARSE_METHODS), ) +GLM4_MOE_LITE_COMPATIBILITY = ModelRuntimeCompatibility( + sparse_methods=frozenset({""}), + prefix_cache_methods=frozenset(), + requires_eager=True, +) + QWEN3_MOE_EP_COMPATIBILITY = ModelRuntimeCompatibility( sparse_methods=_MOE_SPARSE_METHODS, @@ -142,6 +148,7 @@ class ModelRuntimeCompatibility: ("qwen3_5_moe", ParallelMode.OUTER_TP_MOE): QWEN35_MOE_COMPATIBILITY, ("minimax_m2", ParallelMode.STANDARD): MINIMAX_M2_EP_COMPATIBILITY, ("minimax_m2", ParallelMode.OUTER_TP_MOE): MINIMAX_M2_TP_EP_COMPATIBILITY, + ("glm4_moe_lite", ParallelMode.STANDARD): GLM4_MOE_LITE_COMPATIBILITY, } # All shipped cache managers now expose a graph-stable decode preparation path. diff --git a/src/sparsevllm/models/glm4_moe_lite.py b/src/sparsevllm/models/glm4_moe_lite.py new file mode 100644 index 00000000..240f0670 --- /dev/null +++ b/src/sparsevllm/models/glm4_moe_lite.py @@ -0,0 +1,916 @@ +from __future__ import annotations + +import os +import re + +import torch +import torch.nn.functional as F +from torch import nn +from transformers import Glm4MoeLiteConfig + +from sparsevllm.distributed import get_parallel_context +from sparsevllm.engine.cache_manager import MlaLatentWrite +from sparsevllm.layers.embed_head import ParallelLMHead, VocabParallelEmbedding +from sparsevllm.layers.layernorm import RMSNorm +from sparsevllm.layers.linear import ( + ColumnParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from sparsevllm.layers.mla_attention import MLAAttention +from sparsevllm.layers.packed_moe import PackedMoeExperts +from sparsevllm.layers.rotary_embedding import RotaryEmbedding, get_rope +from sparsevllm.models.qwen3 import Qwen3MLP +from sparsevllm.operators.mla_attention import MlaAttentionOpSpec +from sparsevllm.operators.moe import model_activation_dtype, resolve_moe_provider +from sparsevllm.platforms import device_runtime +from sparsevllm.utils.context import get_context +from sparsevllm.utils.log import logger +from sparsevllm.utils.weight_target import WeightTarget + + +_EXPERT_SOURCE_RE = re.compile( + r"^model\.layers\.(\d+)\.mlp\.experts\.(\d+)\." + r"(gate_proj|up_proj|down_proj)\.weight$" +) +_EXPERT_TARGET_RE = re.compile( + r"^model\.layers\.(\d+)\.mlp\.experts\.(\d+)\." + r"(gate_proj|up_proj|down_proj)\.expert_weight$" +) +_MTP_LAYER_INDEX = 47 +_MTP_PREFIX = f"model.layers.{_MTP_LAYER_INDEX}." + + +def _build_mla_attention(config: Glm4MoeLiteConfig) -> MLAAttention: + device = getattr(config, "sparsevllm_device", None) + max_batch_size = int(getattr(config, "sparsevllm_max_batch_size", 0) or 0) + prefill_workspace_bytes = int( + getattr(config, "sparsevllm_mla_prefill_workspace_bytes", 0) or 0 + ) + if device is None or max_batch_size <= 0 or prefill_workspace_bytes <= 0: + raise ValueError( + "GLM MLA construction requires ModelRunner runtime hints: " + "sparsevllm_device, sparsevllm_max_batch_size, and " + "sparsevllm_mla_prefill_workspace_bytes." + ) + parallel_context = get_parallel_context() + activation_dtype = model_activation_dtype(config) + spec = MlaAttentionOpSpec( + num_q_heads=int(config.num_attention_heads), + kv_lora_rank=int(config.kv_lora_rank), + rope_dim=int(config.qk_rope_head_dim), + qk_head_dim=int(config.qk_nope_head_dim + config.qk_rope_head_dim), + value_head_dim=int(config.v_head_dim), + activation_dtype=activation_dtype, + cache_dtype=activation_dtype, + tp_size=int(parallel_context.attention_tp_size), + cuda_graph=bool(getattr(config, "decode_cuda_graph", False)), + ) + return MLAAttention.bind( + spec=spec, + device=device, + max_batch_size=max_batch_size, + prefill_workspace_bytes=prefill_workspace_bytes, + ) + + +class Glm4MoeLiteAttention(nn.Module): + """GLM projections around a shared latent-MLA execution object.""" + + def __init__( + self, + config: Glm4MoeLiteConfig, + mla_attention: MLAAttention, + ) -> None: + super().__init__() + self.mla_attention = mla_attention + self.parallel_context = get_parallel_context() + self.num_heads = int(config.num_attention_heads) + self.local_heads = int(mla_attention.spec.local_q_heads) + self.q_lora_rank = int(config.q_lora_rank) + self.kv_lora_rank = int(config.kv_lora_rank) + self.qk_nope_head_dim = int(config.qk_nope_head_dim) + self.qk_rope_head_dim = int(config.qk_rope_head_dim) + self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + self.v_head_dim = int(config.v_head_dim) + self.proj_chunk_size = int(getattr(config, "mlp_chunk_size", 16384)) + if self.proj_chunk_size <= 0: + raise ValueError( + f"mlp_chunk_size must be positive, got {self.proj_chunk_size}." + ) + quantization = getattr(config, "quantization_config", None) + if bool(getattr(quantization, "enabled", False)): + raise NotImplementedError("GLM MLA projections do not support quantization.") + + self.q_a_proj = ReplicatedLinear( + int(config.hidden_size), + self.q_lora_rank, + bias=bool(config.attention_bias), + ) + self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps) + self.q_b_proj = ColumnParallelLinear( + self.q_lora_rank, + self.num_heads * self.qk_head_dim, + bias=False, + ) + self.kv_a_proj_with_mqa = ReplicatedLinear( + int(config.hidden_size), + self.kv_lora_rank + self.qk_rope_head_dim, + bias=bool(config.attention_bias), + ) + self.kv_a_layernorm = RMSNorm( + self.kv_lora_rank, + eps=config.rms_norm_eps, + ) + self.kv_b_proj = ColumnParallelLinear( + self.kv_lora_rank, + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), + bias=False, + ) + self.o_proj = RowParallelLinear( + self.num_heads * self.v_head_dim, + int(config.hidden_size), + bias=bool(config.attention_bias), + ) + + def _project_kv_history(self, latent: torch.Tensor) -> torch.Tensor: + if int(latent.shape[0]) <= self.proj_chunk_size: + return self.kv_b_proj(latent) + output = torch.empty( + latent.shape[0], + self.local_heads * (self.qk_nope_head_dim + self.v_head_dim), + dtype=latent.dtype, + device=latent.device, + ) + for start in range(0, int(latent.shape[0]), self.proj_chunk_size): + end = min(start + self.proj_chunk_size, int(latent.shape[0])) + output[start:end].copy_(self.kv_b_proj(latent[start:end])) + return output + + def _project_output(self, value_output: torch.Tensor, out: torch.Tensor) -> torch.Tensor: + flattened = value_output.flatten(1, -1) + if int(flattened.shape[0]) <= self.proj_chunk_size: + out.copy_(self.o_proj(flattened)) + return out + for start in range(0, int(flattened.shape[0]), self.proj_chunk_size): + end = min(start + self.proj_chunk_size, int(flattened.shape[0])) + out[start:end].copy_(self.o_proj(flattened[start:end])) + return out + + def _decode_absorbed_query(self, q_nope: torch.Tensor) -> torch.Tensor: + kv_b_weight = self.kv_b_proj.weight.view( + self.local_heads, + self.qk_nope_head_dim + self.v_head_dim, + self.kv_lora_rank, + ) + k_weight = kv_b_weight[:, : self.qk_nope_head_dim] + return torch.bmm( + q_nope.transpose(0, 1), + k_weight, + ).transpose(0, 1) + + def _reconstruct_decode_values(self, latent_output: torch.Tensor) -> torch.Tensor: + kv_b_weight = self.kv_b_proj.weight.view( + self.local_heads, + self.qk_nope_head_dim + self.v_head_dim, + self.kv_lora_rank, + ) + v_weight = kv_b_weight[:, self.qk_nope_head_dim :] + return torch.bmm( + latent_output.transpose(0, 1), + v_weight.transpose(1, 2), + ).transpose(0, 1) + + def _run_attention( + self, + q: torch.Tensor, + q_nope: torch.Tensor, + q_rope: torch.Tensor, + latent: torch.Tensor, + rope: torch.Tensor, + ) -> torch.Tensor: + context = get_context() + cache_manager = context.cache_manager + sparse_controller = context.sparse_controller + layer_idx = int(context.now_layer_idx) + slot_mapping = cache_manager.store_attention_payload( + layer_idx, + MlaLatentWrite( + latent=latent.unsqueeze(1), + rope=rope.unsqueeze(1), + ), + ) + cache_manager.on_kv_stored(layer_idx, latent, slot_mapping) + + temp_slots = None + try: + if context.is_prefill: + selection = sparse_controller.get_prefill_selection(layer_idx) + cache_manager.before_prefill_layer_attention(layer_idx, selection) + view = cache_manager.build_prefill_compute_view( + layer_idx, + latent, + rope, + selection, + ) + temp_slots = view.meta.temp_slots + if context.cu_seqlens_q is None or context.cu_seqlens_q.numel() <= 1: + return torch.empty_like(q) + history = self.mla_attention.prepare_prefill_history(view) + expanded = self._project_kv_history(history.gathered_latent).view( + history.visible_tokens, + self.local_heads, + self.qk_nope_head_dim + self.v_head_dim, + ) + expanded_k_nope, expanded_v = expanded.split( + [self.qk_nope_head_dim, self.v_head_dim], + dim=-1, + ) + expanded_rope = history.gathered_rope[:, None, :].expand( + -1, + self.local_heads, + -1, + ) + workset = self.mla_attention.bind_prefill_kv( + history, + expanded_k=torch.cat((expanded_k_nope, expanded_rope), dim=-1), + expanded_v=expanded_v, + ) + b_start_loc = context.cu_seqlens_q[:-1] + chunk_lens = context.cu_seqlens_q[1:] - context.cu_seqlens_q[:-1] + output = self.mla_attention.run_prefill( + q, + workset, + b_start_loc=b_start_loc, + chunk_lens=chunk_lens, + ) + cache_manager.collect_prefill_attention_score( + layer_idx, + q, + view, + b_start_loc=b_start_loc, + chunk_lens=chunk_lens, + ) + cache_manager.record_prefill_query( + layer_idx, + q, + view, + b_start_loc=b_start_loc, + chunk_lens=chunk_lens, + ) + else: + selection = sparse_controller.get_decode_selection(layer_idx, q) + view = cache_manager.build_decode_compute_view( + layer_idx, + q, + selection, + num_heads=self.local_heads, + num_kv_heads=1, + ) + latent_output = self.mla_attention.run_decode( + self._decode_absorbed_query(q_nope), + q_rope, + view, + ) + output = self._reconstruct_decode_values(latent_output) + cache_manager.record_decode_query(layer_idx, q) + + sparse_controller.on_layer_attention_end(layer_idx) + cache_manager.on_layer_attention_end(layer_idx) + return output + finally: + if temp_slots is not None and temp_slots.numel() > 0: + cache_manager.release_layer_temp_slots(layer_idx, temp_slots) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + rotary_emb: RotaryEmbedding, + ) -> torch.Tensor: + q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) + q = q.view(-1, self.local_heads, self.qk_head_dim) + q_nope, q_rope = q.split( + [self.qk_nope_head_dim, self.qk_rope_head_dim], + dim=-1, + ) + + compressed_kv = self.kv_a_proj_with_mqa(hidden_states) + latent, k_rope = compressed_kv.split( + [self.kv_lora_rank, self.qk_rope_head_dim], + dim=-1, + ) + latent = self.kv_a_layernorm(latent) + q_rope, k_rope = rotary_emb( + positions, + q_rope, + k_rope.unsqueeze(1), + ) + k_rope = k_rope.squeeze(1) + q = torch.cat((q_nope, q_rope), dim=-1) + value_output = self._run_attention( + q, + q_nope, + q_rope, + latent, + k_rope, + ) + return self._project_output(value_output, hidden_states) + + +class Glm4MoeLiteRouter(nn.Module): + def __init__(self, config: Glm4MoeLiteConfig) -> None: + super().__init__() + self.hidden_size = int(config.hidden_size) + self.num_experts = int(config.n_routed_experts) + self.top_k = int(config.num_experts_per_tok) + self.routed_scaling_factor = float(config.routed_scaling_factor) + self.weight = nn.Parameter( + torch.empty( + self.num_experts, + self.hidden_size, + dtype=torch.float32, + ) + ) + self.e_score_correction_bias = nn.Parameter( + torch.empty(self.num_experts, dtype=torch.float32) + ) + from sparsevllm.triton_kernel.moe_biased_sigmoid import ( + topk_biased_sigmoid, + ) + + self.topk_impl = topk_biased_sigmoid + + def forward( + self, + hidden_states: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + router_logits = F.linear(hidden_states.float(), self.weight) + topk_weights, topk_ids = self.topk_impl( + router_logits, + self.e_score_correction_bias, + top_k=self.top_k, + ) + topk_weights = topk_weights * self.routed_scaling_factor + return router_logits, topk_weights, topk_ids + + +class Glm4MoeLitePackedExperts(PackedMoeExperts): + def __init__(self, config: Glm4MoeLiteConfig) -> None: + super().__init__( + num_experts=int(config.n_routed_experts), + hidden_size=int(config.hidden_size), + intermediate_size=int(config.moe_intermediate_size), + top_k=int(config.num_experts_per_tok), + activation_dtype=model_activation_dtype(config), + fp8_enabled=False, + cuda_graph=bool(getattr(config, "decode_cuda_graph", False)), + routing_method="biased_sigmoid", + model_label="GLM-4.7-Flash", + provider_resolver=resolve_moe_provider, + parallel_context=get_parallel_context(), + ) + + +class Glm4MoeLiteSparseMoeBlock(nn.Module): + def __init__(self, config: Glm4MoeLiteConfig) -> None: + super().__init__() + self.parallel_context = get_parallel_context() + self.mlp_chunk_size = int(getattr(config, "mlp_chunk_size", 16384)) + if self.mlp_chunk_size <= 0: + raise ValueError( + f"mlp_chunk_size must be positive, got {self.mlp_chunk_size}." + ) + self.gate = Glm4MoeLiteRouter(config) + self.experts = Glm4MoeLitePackedExperts(config) + self.shared_experts = Qwen3MLP( + hidden_size=int(config.hidden_size), + intermediate_size=( + int(config.moe_intermediate_size) * int(config.n_shared_experts) + ), + hidden_act=str(config.hidden_act), + mlp_chunk_size=self.mlp_chunk_size, + quantization=None, + ) + + def _routed_chunk(self, hidden_states: torch.Tensor) -> torch.Tensor: + _, topk_weights, topk_ids = self.gate(hidden_states) + return self.experts(hidden_states, topk_ids, topk_weights) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if hidden_states.ndim != 2: + raise ValueError( + "Glm4MoeLiteSparseMoeBlock expects [tokens, hidden], got " + f"{tuple(hidden_states.shape)}." + ) + if int(hidden_states.shape[0]) <= self.mlp_chunk_size: + routed = self._routed_chunk(hidden_states) + else: + routed = torch.cat( + [ + self._routed_chunk(chunk) + for chunk in hidden_states.split(self.mlp_chunk_size, dim=0) + ], + dim=0, + ) + self.parallel_context.world_all_reduce(routed) + return routed + self.shared_experts(hidden_states) + + +class Glm4MoeLiteDecoderLayer(nn.Module): + def __init__( + self, + config: Glm4MoeLiteConfig, + layer_idx: int, + mla_attention: MLAAttention, + ) -> None: + super().__init__() + self.self_attn = Glm4MoeLiteAttention(config, mla_attention) + layer_types = list(config.mlp_layer_types) + if len(layer_types) != int(config.num_hidden_layers): + raise ValueError( + "GLM mlp_layer_types length must match num_hidden_layers, got " + f"{len(layer_types)} and {config.num_hidden_layers}." + ) + layer_type = str(layer_types[int(layer_idx)]) + if layer_type == "dense": + self.mlp = Qwen3MLP( + hidden_size=int(config.hidden_size), + intermediate_size=int(config.intermediate_size), + hidden_act=str(config.hidden_act), + mlp_chunk_size=int(getattr(config, "mlp_chunk_size", 16384)), + quantization=None, + ) + elif layer_type == "sparse": + self.mlp = Glm4MoeLiteSparseMoeBlock(config) + else: + raise ValueError( + f"Unsupported GLM MLP layer type at layer {layer_idx}: {layer_type!r}." + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, + eps=config.rms_norm_eps, + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + rotary_emb: RotaryEmbedding, + ) -> tuple[torch.Tensor, torch.Tensor]: + if residual is None: + hidden_states, residual = self.input_layernorm(hidden_states), hidden_states + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + hidden_states = self.self_attn(positions, hidden_states, rotary_emb) + hidden_states, residual = self.post_attention_layernorm( + hidden_states, + residual, + ) + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + +class Glm4MoeLiteModel(nn.Module): + def __init__( + self, + config: Glm4MoeLiteConfig, + mla_attention: MLAAttention, + ) -> None: + super().__init__() + self.config = config + self.mla_attention = mla_attention + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + ) + rope_parameters = getattr(config, "rope_parameters", None) or {} + self.rotary_emb = get_rope( + int(config.qk_rope_head_dim), + rotary_dim=int(config.qk_rope_head_dim), + max_position=int(config.max_position_embeddings), + base=float(rope_parameters.get("rope_theta", 1_000_000.0)), + rope_scaling=None, + backend="torch", + interleaved=True, + ) + self.layers = nn.ModuleList( + [ + Glm4MoeLiteDecoderLayer(config, layer_idx, mla_attention) + for layer_idx in range(int(config.num_hidden_layers)) + ] + ) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.sparse_controller = None + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + ) -> torch.Tensor: + hidden_states = self.embed_tokens(input_ids) + residual = None + context = get_context() + debug_layers_env = os.getenv("SPARSEVLLM_DEBUG_HIDDEN_LAYERS") + debug_layers = None + if debug_layers_env: + debug_layers = { + int(part) for part in debug_layers_env.split(",") if part.strip() + } + self.debug_last_hidden_states = { + -1: hidden_states[-1:].detach().clone(), + } + + for layer_idx, layer in enumerate(self.layers): + context.now_layer_idx = layer_idx + hidden_states, residual = layer( + positions, + hidden_states, + residual, + self.rotary_emb, + ) + if self.sparse_controller is not None: + hidden_states, residual = self.sparse_controller.apply_activation_hook( + layer_idx, + hidden_states, + residual, + context, + ) + if debug_layers is not None and layer_idx in debug_layers: + layer_output = ( + hidden_states if residual is None else hidden_states + residual + ) + self.debug_last_hidden_states[layer_idx] = ( + layer_output[-1:].detach().clone() + ) + if self.sparse_controller is not None: + self.sparse_controller.on_layer_end(layer_idx, context) + + hidden_states, _ = self.norm(hidden_states, residual) + if debug_layers is not None: + self.debug_last_hidden_states[int(self.config.num_hidden_layers)] = ( + hidden_states[-1:].detach().clone() + ) + return hidden_states + + +def _expected_mtp_weight_names(num_experts: int) -> set[str]: + prefix = _MTP_PREFIX + names = { + prefix + suffix + for suffix in ( + "eh_proj.weight", + "embed_tokens.weight", + "enorm.weight", + "hnorm.weight", + "input_layernorm.weight", + "post_attention_layernorm.weight", + "self_attn.kv_a_layernorm.weight", + "self_attn.kv_a_proj_with_mqa.weight", + "self_attn.kv_b_proj.weight", + "self_attn.o_proj.weight", + "self_attn.q_a_layernorm.weight", + "self_attn.q_a_proj.weight", + "self_attn.q_b_proj.weight", + "mlp.gate.e_score_correction_bias", + "mlp.gate.weight", + "mlp.shared_experts.down_proj.weight", + "mlp.shared_experts.gate_proj.weight", + "mlp.shared_experts.up_proj.weight", + "shared_head.head.weight", + "shared_head.norm.weight", + ) + } + names.update( + f"{prefix}mlp.experts.{expert_id}.{projection}.weight" + for expert_id in range(int(num_experts)) + for projection in ("gate_proj", "up_proj", "down_proj") + ) + return names + + +class Glm4MoeLiteForCausalLM(nn.Module): + special_weight_loaders = (".expert_weight",) + packed_modules_mapping = { + "gate_proj": ("gate_up_proj", 0), + "up_proj": ("gate_up_proj", 1), + } + + def __init__( + self, + config: Glm4MoeLiteConfig, + *, + mla_attention: MLAAttention | None = None, + ) -> None: + super().__init__() + self.config = config + self.parallel_context = get_parallel_context() + if mla_attention is None: + mla_attention = _build_mla_attention(config) + self.model = Glm4MoeLiteModel(config, mla_attention) + self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size) + if config.tie_word_embeddings: + self.lm_head.weight.data = self.model.embed_tokens.weight.data + self.expect_mtp_weights = bool( + getattr(config, "sparsevllm_expect_mtp_weights", False) + ) + self._intentionally_skipped_expert_weights: set[str] = set() + self._intentionally_skipped_mtp_weights: set[str] = set() + + def _sparse_block(self, layer_idx: int) -> Glm4MoeLiteSparseMoeBlock: + if not 0 <= int(layer_idx) < len(self.model.layers): + raise ValueError( + f"GLM expert checkpoint layer {layer_idx} is outside the base model." + ) + block = self.model.layers[int(layer_idx)].mlp + if not isinstance(block, Glm4MoeLiteSparseMoeBlock): + raise ValueError( + f"GLM checkpoint contains experts for dense layer {layer_idx}." + ) + return block + + def iter_tiny_reference_weights( + self, + state_dict: dict[str, torch.Tensor], + ): + """Expand Transformers' packed tiny experts into checkpoint-style names.""" + + gate_up_suffix = ".mlp.experts.gate_up_proj" + down_suffix = ".mlp.experts.down_proj" + for source_name, weight in state_dict.items(): + if source_name.endswith(gate_up_suffix): + prefix = source_name[: -len("gate_up_proj")] + intermediate_size = int(weight.shape[1]) // 2 + for expert_id in range(int(weight.shape[0])): + yield ( + f"{prefix}{expert_id}.gate_proj.weight", + weight[expert_id, :intermediate_size], + ) + yield ( + f"{prefix}{expert_id}.up_proj.weight", + weight[expert_id, intermediate_size:], + ) + continue + if source_name.endswith(down_suffix): + prefix = source_name[: -len("down_proj")] + for expert_id in range(int(weight.shape[0])): + yield ( + f"{prefix}{expert_id}.down_proj.weight", + weight[expert_id], + ) + continue + yield source_name, weight + + @torch.inference_mode() + def warmup_moe(self, num_tokens: int = 1) -> None: + num_tokens = int(num_tokens) + if num_tokens <= 0: + raise ValueError(f"num_tokens must be positive, got {num_tokens}.") + block = next( + ( + layer.mlp + for layer in self.model.layers + if isinstance(layer.mlp, Glm4MoeLiteSparseMoeBlock) + ), + None, + ) + if block is None: + raise RuntimeError("GLM model has no sparse MoE layer to warm up.") + experts = block.experts + hidden_states = torch.zeros( + (num_tokens, experts.hidden_size), + dtype=model_activation_dtype(self.config), + device=experts.w13_weight.device, + ) + block.gate(hidden_states) + top_k = int(self.config.num_experts_per_tok) + topk_ids = ( + torch.arange( + num_tokens * top_k, + dtype=torch.int64, + device=hidden_states.device, + ) + .remainder(experts.num_local_experts) + .add(experts.local_expert_start) + .view(num_tokens, top_k) + ) + topk_weights = torch.full( + (num_tokens, top_k), + 1.0 / top_k, + dtype=torch.float32, + device=hidden_states.device, + ) + experts(hidden_states, topk_ids, topk_weights) + device_runtime.synchronize() + + def map_weight_name(self, source_weight_name: str) -> str | None: + if source_weight_name.startswith(_MTP_PREFIX): + return None + match = _EXPERT_SOURCE_RE.match(source_weight_name) + if match is None: + return source_weight_name + layer_idx, global_expert_id, projection = match.groups() + block = self._sparse_block(int(layer_idx)) + global_expert_id = int(global_expert_id) + if not block.experts.is_local_expert(global_expert_id): + self._intentionally_skipped_expert_weights.add(source_weight_name) + return None + return ( + f"model.layers.{layer_idx}.mlp.experts.{global_expert_id}." + f"{projection}.expert_weight" + ) + + def resolve_special_weight( + self, + target_weight_name: str, + ) -> WeightTarget | None: + match = _EXPERT_TARGET_RE.match(target_weight_name) + if match is None: + return None + layer_idx, global_expert_id, projection = match.groups() + experts = self._sparse_block(int(layer_idx)).experts + return WeightTarget(experts, (int(global_expert_id), projection)) + + def record_skipped_weight( + self, + source_weight_name: str, + loaded_weight_shape: tuple[int, ...] | None, + loaded_weight_dtype: str | None, + loaded_scale_shape: tuple[int, ...] | None, + loaded_scale_dtype: str | None, + ) -> None: + if source_weight_name.startswith(_MTP_PREFIX): + if loaded_weight_shape is None or any( + int(dim) <= 0 for dim in loaded_weight_shape + ): + raise ValueError( + f"Invalid GLM MTP tensor shape for {source_weight_name!r}: " + f"{loaded_weight_shape}." + ) + if loaded_weight_dtype not in {"BF16", "F32"}: + raise TypeError( + f"Unsupported GLM MTP dtype for {source_weight_name!r}: " + f"{loaded_weight_dtype}." + ) + if loaded_scale_shape is not None or loaded_scale_dtype is not None: + raise ValueError( + f"Unquantized GLM MTP tensor has an unexpected scale: " + f"{source_weight_name!r}." + ) + self._intentionally_skipped_mtp_weights.add(source_weight_name) + return + + match = _EXPERT_SOURCE_RE.match(source_weight_name) + if match is None: + raise ValueError( + f"GLM loader unexpectedly skipped {source_weight_name!r}." + ) + layer_idx, global_expert_id, projection = match.groups() + experts = self._sparse_block(int(layer_idx)).experts + global_expert_id = int(global_expert_id) + if experts.is_local_expert(global_expert_id): + raise ValueError( + f"GLM loader skipped local expert weight {source_weight_name!r}." + ) + expected_shape = ( + (experts.hidden_size, experts.global_intermediate_size) + if projection == "down_proj" + else (experts.global_intermediate_size, experts.hidden_size) + ) + if loaded_weight_shape != expected_shape: + raise ValueError( + "Remote GLM expert weight shape mismatch for " + f"{source_weight_name!r}: expected={expected_shape}, " + f"got={loaded_weight_shape}." + ) + if loaded_weight_dtype not in {"BF16", "F16", "F32"}: + raise TypeError( + f"Unsupported remote GLM expert dtype for " + f"{source_weight_name!r}: {loaded_weight_dtype}." + ) + if loaded_scale_shape is not None or loaded_scale_dtype is not None: + raise ValueError( + f"Unquantized remote GLM expert has an unexpected scale: " + f"{source_weight_name!r}." + ) + self._intentionally_skipped_expert_weights.add(source_weight_name) + + def load_special_weight( + self, + target_weight_name: str, + loaded_weight: torch.Tensor, + loaded_scale: torch.Tensor | None, + ) -> int: + target = self.resolve_special_weight(target_weight_name) + if target is None: + return 0 + global_expert_id, projection = target.shard_id + target.module.load_expert_weight( + global_expert_id, + projection, + loaded_weight, + loaded_scale, + ) + return 1 + + def validate_loaded_weights(self, loaded_parameter_names: set[str]) -> None: + packed_expert_parameters = { + name + for name, _ in self.named_parameters() + if name.endswith(".mlp.experts.w13_weight") + or name.endswith(".mlp.experts.w2_weight") + } + expected_dense_parameters = { + name for name, _ in self.named_parameters() + } - packed_expert_parameters + missing_dense = sorted(expected_dense_parameters - loaded_parameter_names) + if missing_dense: + raise ValueError( + f"Missing replicated GLM base weights: {missing_dense[:8]}." + ) + + sparse_blocks = [ + layer.mlp + for layer in self.model.layers + if isinstance(layer.mlp, Glm4MoeLiteSparseMoeBlock) + ] + for block in sparse_blocks: + block.experts.validate_loaded_weights() + + expected_remote = { + f"model.layers.{layer_idx}.mlp.experts.{expert_id}." + f"{projection}.weight" + for layer_idx, layer in enumerate(self.model.layers) + if isinstance(layer.mlp, Glm4MoeLiteSparseMoeBlock) + for expert_id in range(int(self.config.n_routed_experts)) + if not layer.mlp.experts.is_local_expert(expert_id) + for projection in ("gate_proj", "up_proj", "down_proj") + } + missing_remote = sorted( + expected_remote - self._intentionally_skipped_expert_weights + ) + unexpected_remote = sorted( + self._intentionally_skipped_expert_weights - expected_remote + ) + if missing_remote or unexpected_remote: + raise ValueError( + "GLM remote expert skip set is inconsistent: " + f"missing={missing_remote[:4]} unexpected={unexpected_remote[:4]}." + ) + + expected_mtp = _expected_mtp_weight_names(self.config.n_routed_experts) + if self.expect_mtp_weights: + missing_mtp = sorted( + expected_mtp - self._intentionally_skipped_mtp_weights + ) + unexpected_mtp = sorted( + self._intentionally_skipped_mtp_weights - expected_mtp + ) + if missing_mtp or unexpected_mtp: + raise ValueError( + "GLM MTP skip set is inconsistent: " + f"missing={missing_mtp[:4]} unexpected={unexpected_mtp[:4]}." + ) + elif self._intentionally_skipped_mtp_weights: + raise ValueError( + "GLM checkpoint supplied MTP weights when they were not expected: " + f"{sorted(self._intentionally_skipped_mtp_weights)[:4]}." + ) + + first_sparse = sparse_blocks[0] if sparse_blocks else None + logger.info( + "Loaded GLM-4.7-Flash rank {} MLA provider={} MoE provider={} " + "attention TP {}/{} across {} base layers; local expert shards={} " + "remote expert skips={} MTP skips={}.", + self.parallel_context.world_rank, + self.model.mla_attention.provider.name, + first_sparse.experts.provider.name if first_sparse is not None else "none", + self.parallel_context.attention_tp_rank, + self.parallel_context.attention_tp_size, + len(self.model.layers), + sum(len(block.experts._loaded_expert_shards) for block in sparse_blocks), + len(self._intentionally_skipped_expert_weights), + len(self._intentionally_skipped_mtp_weights), + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + ) -> torch.Tensor: + return self.model(input_ids, positions) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.lm_head(hidden_states) + + +__all__ = [ + "Glm4MoeLiteAttention", + "Glm4MoeLiteForCausalLM", + "Glm4MoeLiteModel", + "Glm4MoeLitePackedExperts", + "Glm4MoeLiteRouter", + "Glm4MoeLiteSparseMoeBlock", +] diff --git a/src/sparsevllm/models/spec.py b/src/sparsevllm/models/spec.py index 439ef192..b9b41c72 100644 --- a/src/sparsevllm/models/spec.py +++ b/src/sparsevllm/models/spec.py @@ -159,6 +159,14 @@ def validate_sharding(self, hf_config: Any, topology: ParallelTopology) -> None: moe_tp_fields=("intermediate_size",), top_k_field="num_experts_per_tok", ), + "glm4_moe_lite": ModelSpec( + "GLM-4.7-Flash", + runtime_class_name="Glm4MoeLiteForCausalLM", + attention_tp_fields=("num_attention_heads", "vocab_size"), + num_experts_field="n_routed_experts", + moe_tp_fields=("intermediate_size", "moe_intermediate_size"), + top_k_field="num_experts_per_tok", + ), } ) diff --git a/src/sparsevllm/triton_kernel/minimax_m2_router.py b/src/sparsevllm/triton_kernel/minimax_m2_router.py index d0f7513f..48c9004a 100644 --- a/src/sparsevllm/triton_kernel/minimax_m2_router.py +++ b/src/sparsevllm/triton_kernel/minimax_m2_router.py @@ -2,46 +2,16 @@ import torch import torch.nn.functional as F -import triton -import triton.language as tl + +from sparsevllm.triton_kernel.moe_biased_sigmoid import ( + topk_biased_sigmoid as _topk_biased_sigmoid, +) _NUM_EXPERTS = 256 _TOP_K = 8 -@triton.jit -def _topk_biased_sigmoid_kernel( - routing_weights_ptr, - correction_bias_ptr, - ids_ptr, - stride_routing_weights_m, - stride_ids_m, -): - row = tl.program_id(0) - offsets = tl.arange(0, 256) - routing_weights = tl.load( - routing_weights_ptr + row * stride_routing_weights_m + offsets - ) - scores = routing_weights + tl.load(correction_bias_ptr + offsets) - - # CUDA topk(sorted=False) writes values strictly above the kth threshold - # first, then fills the remaining slots with first-seen threshold ties. - selection_values = tl.where(scores == scores, scores, float("inf")) - threshold = tl.min(tl.topk(selection_values, 8), axis=0) - greater_mask = selection_values > threshold - equal_mask = selection_values == threshold - greater_rank = tl.cumsum(greater_mask.to(tl.int32), axis=0) - 1 - equal_rank = tl.cumsum(equal_mask.to(tl.int32), axis=0) - 1 - num_greater = tl.sum(greater_mask.to(tl.int32), axis=0) - selected_equal = equal_mask & (equal_rank < 8 - num_greater) - selected = greater_mask | selected_equal - output_slot = tl.where(greater_mask, greater_rank, num_greater + equal_rank) - - ids_base = ids_ptr + row * stride_ids_m - tl.store(ids_base + output_slot, offsets, mask=selected) - - def _validate_router_inputs( router_logits: torch.Tensor, correction_bias: torch.Tensor, @@ -85,22 +55,12 @@ def topk_biased_sigmoid( ) -> tuple[torch.Tensor, torch.Tensor]: """Route MiniMax M2.7 logits with its biased-sigmoid top-k rule.""" _validate_router_inputs(router_logits, correction_bias, top_k) - num_tokens = int(router_logits.shape[0]) - routing_weights = torch.sigmoid(router_logits) - ids = torch.empty( - (num_tokens, _TOP_K), dtype=torch.int64, device=router_logits.device - ) - _topk_biased_sigmoid_kernel[(num_tokens,)]( - routing_weights, + return _topk_biased_sigmoid( + router_logits, correction_bias, - ids, - routing_weights.stride(0), - ids.stride(0), - num_warps=2 if num_tokens <= 256 else 1, + top_k=top_k, + normalization_epsilon=0.0, ) - weights = routing_weights.gather(1, ids) - weights = weights / weights.sum(dim=-1, keepdim=True) - return weights, ids def minimax_m2_router( diff --git a/src/sparsevllm/triton_kernel/moe_biased_sigmoid.py b/src/sparsevllm/triton_kernel/moe_biased_sigmoid.py new file mode 100644 index 00000000..00ef78d5 --- /dev/null +++ b/src/sparsevllm/triton_kernel/moe_biased_sigmoid.py @@ -0,0 +1,133 @@ +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +_SUPPORTED_SHAPES = frozenset({(64, 4), (256, 8)}) + + +@triton.jit +def _topk_biased_sigmoid_kernel( + routing_weights_ptr, + correction_bias_ptr, + ids_ptr, + stride_routing_weights_m, + stride_ids_m, + NUM_EXPERTS: tl.constexpr, + TOP_K: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + row = tl.program_id(0) + offsets = tl.arange(0, BLOCK_SIZE) + expert_mask = offsets < NUM_EXPERTS + routing_weights = tl.load( + routing_weights_ptr + row * stride_routing_weights_m + offsets, + mask=expert_mask, + other=-float("inf"), + ) + correction_bias = tl.load( + correction_bias_ptr + offsets, + mask=expert_mask, + other=0.0, + ) + scores = routing_weights + correction_bias + + # Match torch.topk(sorted=False): values strictly above the kth threshold + # are emitted first, followed by first-seen threshold ties. + selection_values = tl.where(scores == scores, scores, float("inf")) + threshold = tl.min(tl.topk(selection_values, TOP_K), axis=0) + greater_mask = expert_mask & (selection_values > threshold) + equal_mask = expert_mask & (selection_values == threshold) + greater_rank = tl.cumsum(greater_mask.to(tl.int32), axis=0) - 1 + equal_rank = tl.cumsum(equal_mask.to(tl.int32), axis=0) - 1 + num_greater = tl.sum(greater_mask.to(tl.int32), axis=0) + selected_equal = equal_mask & (equal_rank < TOP_K - num_greater) + selected = greater_mask | selected_equal + output_slot = tl.where(greater_mask, greater_rank, num_greater + equal_rank) + + ids_base = ids_ptr + row * stride_ids_m + tl.store(ids_base + output_slot, offsets, mask=selected) + + +def _validate_inputs( + router_logits: torch.Tensor, + correction_bias: torch.Tensor, + *, + top_k: int, +) -> tuple[int, int]: + if not router_logits.is_cuda or not correction_bias.is_cuda: + raise ValueError("Biased-sigmoid routing requires CUDA tensors.") + if router_logits.device != correction_bias.device: + raise ValueError("router_logits and correction_bias must be on one device.") + if router_logits.ndim != 2: + raise ValueError( + "router_logits must have shape [tokens, experts], got " + f"{tuple(router_logits.shape)}." + ) + num_tokens, num_experts = map(int, router_logits.shape) + if tuple(correction_bias.shape) != (num_experts,): + raise ValueError( + f"correction_bias must have shape [{num_experts}], got " + f"{tuple(correction_bias.shape)}." + ) + if router_logits.dtype != torch.float32 or correction_bias.dtype != torch.float32: + raise TypeError( + "Biased-sigmoid routing requires FP32 logits and correction_bias, " + f"got {router_logits.dtype} and {correction_bias.dtype}." + ) + if not router_logits.is_contiguous() or not correction_bias.is_contiguous(): + raise ValueError("Biased-sigmoid router inputs must be contiguous.") + if num_tokens <= 0: + raise ValueError("Biased-sigmoid routing requires at least one token.") + shape = (num_experts, int(top_k)) + if shape not in _SUPPORTED_SHAPES: + raise ValueError( + "Unsupported biased-sigmoid router shape: " + f"num_experts={num_experts}, top_k={top_k}; supported=" + f"{sorted(_SUPPORTED_SHAPES)}." + ) + return num_tokens, num_experts + + +def topk_biased_sigmoid( + router_logits: torch.Tensor, + correction_bias: torch.Tensor, + *, + top_k: int, + normalization_epsilon: float = 1e-20, +) -> tuple[torch.Tensor, torch.Tensor]: + """Select experts with correction bias and return unbiased sigmoid weights.""" + + num_tokens, num_experts = _validate_inputs( + router_logits, + correction_bias, + top_k=top_k, + ) + routing_weights = torch.sigmoid(router_logits) + ids = torch.empty( + (num_tokens, int(top_k)), + dtype=torch.int64, + device=router_logits.device, + ) + block_size = triton.next_power_of_2(num_experts) + _topk_biased_sigmoid_kernel[(num_tokens,)]( + routing_weights, + correction_bias, + ids, + routing_weights.stride(0), + ids.stride(0), + NUM_EXPERTS=num_experts, + TOP_K=int(top_k), + BLOCK_SIZE=block_size, + num_warps=2 if num_tokens <= 256 else 1, + ) + weights = routing_weights.gather(1, ids) + weights = weights / ( + weights.sum(dim=-1, keepdim=True) + float(normalization_epsilon) + ) + return weights, ids + + +__all__ = ["topk_biased_sigmoid"] diff --git a/tests/test_attention_cache_storage.py b/tests/test_attention_cache_storage.py index 8babfec9..64fafcde 100644 --- a/tests/test_attention_cache_storage.py +++ b/tests/test_attention_cache_storage.py @@ -33,12 +33,29 @@ def _glm_hf_config(**overrides): "max_position_embeddings": 4096, "hidden_size": 64, "intermediate_size": 128, - "num_hidden_layers": 2, - "num_attention_heads": 4, - "num_key_value_heads": 4, + "num_hidden_layers": 47, + "num_attention_heads": 20, + "num_key_value_heads": 20, "vocab_size": 128, + "q_lora_rank": 768, "kv_lora_rank": 512, + "qk_nope_head_dim": 192, "qk_rope_head_dim": 64, + "v_head_dim": 256, + "moe_intermediate_size": 64, + "n_routed_experts": 64, + "n_shared_experts": 1, + "num_experts_per_tok": 4, + "n_group": 1, + "topk_group": 1, + "topk_method": "noaux_tc", + "norm_topk_prob": True, + "routed_scaling_factor": 1.8, + "mlp_layer_types": ["dense"] + ["sparse"] * 46, + "num_nextn_predict_layers": 1, + "rope_interleave": True, + "rope_scaling": None, + "attention_bias": False, "quantization_config": None, } values.update(overrides) diff --git a/tests/test_glm4_moe_lite.py b/tests/test_glm4_moe_lite.py new file mode 100644 index 00000000..1e913ade --- /dev/null +++ b/tests/test_glm4_moe_lite.py @@ -0,0 +1,322 @@ +from __future__ import annotations + +from contextlib import ExitStack, contextmanager +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch +import torch.nn.functional as F +from transformers import Glm4MoeLiteConfig +from transformers.models.glm4_moe_lite.modeling_glm4_moe_lite import ( + apply_rotary_pos_emb_interleave, +) + +from sparsevllm.config import QuantizationConfig +from sparsevllm.debug.tiny_random import ( + build_tiny_random_hf_model, + initialize_sparse_model, +) +from sparsevllm.distributed import ParallelContext, ParallelGroup +from sparsevllm.layers.rotary_embedding import apply_interleaved_rotary_emb +from sparsevllm.models.glm4_moe_lite import ( + Glm4MoeLiteAttention, + Glm4MoeLiteForCausalLM, + Glm4MoeLiteRouter, + Glm4MoeLiteSparseMoeBlock, + _expected_mtp_weight_names, +) +from sparsevllm.models.qwen3 import Qwen3MLP +from sparsevllm.operators.mla_attention import MlaAttentionOpSpec +from sparsevllm.operators.moe import TritonMoeProvider + + +def _config(**overrides) -> Glm4MoeLiteConfig: + values = { + "vocab_size": 128, + "hidden_size": 64, + "intermediate_size": 128, + "moe_intermediate_size": 16, + "num_hidden_layers": 2, + "num_attention_heads": 20, + "num_key_value_heads": 20, + "n_shared_experts": 1, + "n_routed_experts": 64, + "num_experts_per_tok": 4, + "routed_scaling_factor": 1.8, + "n_group": 1, + "topk_group": 1, + "norm_topk_prob": True, + "kv_lora_rank": 512, + "q_lora_rank": 768, + "qk_rope_head_dim": 64, + "v_head_dim": 256, + "qk_nope_head_dim": 192, + "max_position_embeddings": 32, + "dtype": torch.bfloat16, + "rope_parameters": { + "rope_type": "default", + "rope_theta": 1_000_000.0, + }, + } + values.update(overrides) + config = Glm4MoeLiteConfig(**values) + config.mlp_chunk_size = 8 + config.quantization_config = QuantizationConfig.disabled() + config.decode_cuda_graph = False + config.sparsevllm_expect_mtp_weights = False + return config + + +def _tp_context(tp_rank: int = 0, tp_size: int = 1) -> ParallelContext: + ranks = tuple(range(tp_size)) + return ParallelContext( + world=ParallelGroup(None, ranks, tp_rank, tp_size), + tensor=ParallelGroup(None, ranks, tp_rank, tp_size), + expert=ParallelGroup(None, (tp_rank,), 0, 1), + data=ParallelGroup(None, (tp_rank,), 0, 1), + ) + + +def _fake_mla(tp_size: int = 1): + spec = MlaAttentionOpSpec( + num_q_heads=20, + kv_lora_rank=512, + rope_dim=64, + qk_head_dim=256, + value_head_dim=256, + activation_dtype=torch.bfloat16, + cache_dtype=torch.bfloat16, + tp_size=tp_size, + cuda_graph=False, + ) + return SimpleNamespace( + spec=spec, + provider=SimpleNamespace(name="test"), + ) + + +@contextmanager +def _construction_context(context: ParallelContext): + with ExitStack() as stack: + stack.enter_context( + patch( + "sparsevllm.models.glm4_moe_lite.get_parallel_context", + return_value=context, + ) + ) + stack.enter_context( + patch("sparsevllm.layers.linear.get_parallel_context", return_value=context) + ) + stack.enter_context( + patch( + "sparsevllm.layers.embed_head.get_parallel_context", + return_value=context, + ) + ) + stack.enter_context( + patch( + "sparsevllm.models.glm4_moe_lite.resolve_moe_provider", + return_value=TritonMoeProvider(), + ) + ) + stack.enter_context(torch.device("cpu")) + previous_dtype = torch.get_default_dtype() + torch.set_default_dtype(torch.bfloat16) + try: + yield + finally: + torch.set_default_dtype(previous_dtype) + + +def _model(config=None, *, tp_rank: int = 0, tp_size: int = 1): + config = _config() if config is None else config + context = _tp_context(tp_rank, tp_size) + with _construction_context(context): + return Glm4MoeLiteForCausalLM( + config, + mla_attention=_fake_mla(tp_size), + ) + + +def test_glm_topology_reuses_one_mla_object_and_qwen_dense_mlp() -> None: + model = _model() + + assert len(model.model.layers) == 2 + assert isinstance(model.model.layers[0].mlp, Qwen3MLP) + assert isinstance(model.model.layers[1].mlp, Glm4MoeLiteSparseMoeBlock) + assert model.model.layers[0].self_attn.mla_attention is model.model.mla_attention + assert model.model.layers[1].self_attn.mla_attention is model.model.mla_attention + assert model.model.rotary_emb.interleaved is True + assert model.model.layers[1].mlp.experts.op_spec.routing_method == "biased_sigmoid" + + +def test_glm_interleaved_rope_matches_transformers() -> None: + torch.manual_seed(13) + q = torch.randn(1, 3, 5, 64) + k = torch.randn(1, 1, 5, 64) + angles = torch.randn(1, 5, 32) + cos_half = angles.cos() + sin_half = angles.sin() + cos = torch.cat((cos_half, cos_half), dim=-1) + sin = torch.cat((sin_half, sin_half), dim=-1) + + expected_q, expected_k = apply_rotary_pos_emb_interleave(q, k, cos, sin) + actual_q = apply_interleaved_rotary_emb( + q, + cos_half.unsqueeze(1), + sin_half.unsqueeze(1), + ) + actual_k = apply_interleaved_rotary_emb( + k, + cos_half.unsqueeze(1), + sin_half.unsqueeze(1), + ) + + torch.testing.assert_close(actual_q, expected_q) + torch.testing.assert_close(actual_k, expected_k) + + +def test_glm_tp_projection_slices_follow_local_heads() -> None: + config = _config() + context = _tp_context(tp_rank=2, tp_size=4) + with _construction_context(context): + attention = Glm4MoeLiteAttention(config, _fake_mla(tp_size=4)) + + q_source = ( + torch.arange(20 * 256 * 768).remainder(127).to(torch.bfloat16) + ).view(20 * 256, 768) + kv_source = ( + torch.arange(20 * 448 * 512).remainder(127).to(torch.bfloat16) + ).view(20 * 448, 512) + o_source = ( + torch.arange(64 * 20 * 256).remainder(127).to(torch.bfloat16) + ).view(64, 20 * 256) + attention.q_b_proj.weight_loader(attention.q_b_proj.weight, q_source) + attention.kv_b_proj.weight_loader(attention.kv_b_proj.weight, kv_source) + attention.o_proj.weight_loader(attention.o_proj.weight, o_source) + + assert attention.local_heads == 5 + assert torch.equal(attention.q_b_proj.weight, q_source[2560:3840]) + assert torch.equal(attention.kv_b_proj.weight, kv_source[4480:6720]) + assert torch.equal(attention.o_proj.weight, o_source[:, 2560:3840]) + + +def test_glm_decode_absorption_and_value_reconstruction_match_linear_algebra() -> None: + config = _config() + context = _tp_context() + with _construction_context(context): + attention = Glm4MoeLiteAttention(config, _fake_mla()) + torch.manual_seed(19) + attention.kv_b_proj.weight.data.normal_(mean=0.0, std=0.02) + q_nope = torch.randn(3, 20, 192, dtype=torch.bfloat16) + latent_output = torch.randn(3, 20, 512, dtype=torch.bfloat16) + weight = attention.kv_b_proj.weight.view(20, 448, 512) + + absorbed = attention._decode_absorbed_query(q_nope) + reconstructed = attention._reconstruct_decode_values(latent_output) + + expected_absorbed = torch.einsum( + "bhd,hdr->bhr", + q_nope, + weight[:, :192], + ) + expected_reconstructed = torch.einsum( + "bhr,hvr->bhv", + latent_output, + weight[:, 192:], + ) + torch.testing.assert_close(absorbed, expected_absorbed) + torch.testing.assert_close(reconstructed, expected_reconstructed) + + +def test_glm_router_uses_bias_only_for_selection_and_scales_weights() -> None: + torch.manual_seed(23) + router = Glm4MoeLiteRouter(_config()) + router.weight.data.normal_(mean=0.0, std=0.1) + router.e_score_correction_bias.data.zero_() + router.e_score_correction_bias.data[:4] = 100.0 + + def torch_topk(logits, correction_bias, *, top_k): + scores = logits.sigmoid() + ids = torch.topk( + scores + correction_bias, + top_k, + dim=-1, + sorted=False, + ).indices + weights = scores.gather(1, ids) + return weights / (weights.sum(dim=-1, keepdim=True) + 1e-20), ids + + router.topk_impl = torch_topk + hidden_states = torch.randn(7, 64, dtype=torch.bfloat16) + logits, weights, ids = router(hidden_states) + original_scores = logits.sigmoid() + expected = original_scores.gather(1, ids) + expected = expected / (expected.sum(dim=-1, keepdim=True) + 1e-20) * 1.8 + + assert all(set(row.tolist()) == {0, 1, 2, 3} for row in ids) + torch.testing.assert_close(weights, expected) + torch.testing.assert_close(weights.sum(dim=-1), torch.full((7,), 1.8)) + + +def test_tiny_transformers_weights_load_through_strict_glm_mapping() -> None: + config = _config() + context = _tp_context() + with _construction_context(context): + model = Glm4MoeLiteForCausalLM(config, mla_attention=_fake_mla()) + initialize_sparse_model(model, config, seed=29) + reference = build_tiny_random_hf_model(config, seed=29) + reference_experts = reference.model.layers[1].mlp.experts + target_experts = model.model.layers[1].mlp.experts + + torch.testing.assert_close( + model.model.layers[0].self_attn.q_b_proj.weight.cpu(), + reference.model.layers[0].self_attn.q_b_proj.weight, + ) + torch.testing.assert_close( + target_experts.w13_weight[:, :16].cpu(), + reference_experts.gate_up_proj[:, :16], + ) + torch.testing.assert_close( + target_experts.w13_weight[:, 16:].cpu(), + reference_experts.gate_up_proj[:, 16:], + ) + torch.testing.assert_close( + target_experts.w2_weight.cpu(), + reference_experts.down_proj, + ) + assert len(target_experts._loaded_expert_shards) == 64 * 3 + assert model.model.layers[1].mlp.gate.e_score_correction_bias.dtype == torch.float32 + + +def test_glm_mtp_skip_set_is_exact() -> None: + config = _config() + config.sparsevllm_expect_mtp_weights = True + model = _model(config) + dense_names = { + name + for name, _ in model.named_parameters() + if not name.endswith(".mlp.experts.w13_weight") + and not name.endswith(".mlp.experts.w2_weight") + } + for block in ( + layer.mlp + for layer in model.model.layers + if isinstance(layer.mlp, Glm4MoeLiteSparseMoeBlock) + ): + block.experts._loaded_expert_shards = { + (expert_id, projection) + for expert_id in range(64) + for projection in ("gate_proj", "up_proj", "down_proj") + } + expected_mtp = _expected_mtp_weight_names(64) + for name in expected_mtp: + model.record_skipped_weight(name, (1,), "BF16", None, None) + + model.validate_loaded_weights(dense_names) + assert len(model._intentionally_skipped_mtp_weights) == 212 + + model._intentionally_skipped_mtp_weights.remove(next(iter(expected_mtp))) + with pytest.raises(ValueError, match="MTP skip set is inconsistent"): + model.validate_loaded_weights(dense_names) diff --git a/tests/test_minimax_m2_router.py b/tests/test_minimax_m2_router.py index 8caba566..8abe58a3 100644 --- a/tests/test_minimax_m2_router.py +++ b/tests/test_minimax_m2_router.py @@ -8,6 +8,9 @@ minimax_m2_router, topk_biased_sigmoid, ) +from sparsevllm.triton_kernel.moe_biased_sigmoid import ( + topk_biased_sigmoid as generic_topk_biased_sigmoid, +) def _reference( @@ -36,6 +39,32 @@ def test_topk_biased_sigmoid_matches_minimax_reference(): assert torch.equal(weights, expected_weights) +@unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for this test.") +def test_topk_biased_sigmoid_matches_glm_64_expert_reference(): + torch.manual_seed(29) + logits = torch.randn(1024, 64, dtype=torch.float32, device="cuda") * 3 + correction_bias = torch.randn(64, dtype=torch.float32, device="cuda") * 0.1 + routing_weights = torch.sigmoid(logits) + expected_ids = torch.topk( + routing_weights + correction_bias, + 4, + dim=-1, + sorted=False, + ).indices + expected_weights = routing_weights.gather(1, expected_ids) + expected_weights /= expected_weights.sum(dim=-1, keepdim=True) + 1e-20 + + weights, ids = generic_topk_biased_sigmoid( + logits, + correction_bias, + top_k=4, + ) + torch.cuda.synchronize() + + assert torch.equal(ids, expected_ids) + assert torch.equal(weights, expected_weights) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for this test.") def test_topk_biased_sigmoid_matches_nonfinite_reference(): logits = torch.zeros(6, 256, dtype=torch.float32, device="cuda") From c947de47398063d6325950bb734807cd851ac01e Mon Sep 17 00:00:00 2001 From: Codex Date: Tue, 4 Aug 2026 23:43:07 +0800 Subject: [PATCH 08/48] feat: add glm response template --- .../entrypoints/openai/api_server.py | 2 +- .../openai/serving/response_parsing.py | 52 +++++++++ tests/test_openai_api_server.py | 105 +++++++++++++++++- 3 files changed, 156 insertions(+), 3 deletions(-) diff --git a/src/sparsevllm/entrypoints/openai/api_server.py b/src/sparsevllm/entrypoints/openai/api_server.py index 3b7bec8e..9d584e39 100644 --- a/src/sparsevllm/entrypoints/openai/api_server.py +++ b/src/sparsevllm/entrypoints/openai/api_server.py @@ -73,7 +73,7 @@ UNSUPPORTED_SERVING_METHOD_PREFIXES = ("deltakv",) -SUPPORTED_RESPONSE_PARSERS = ("qwen3", "minimax_m2") +SUPPORTED_RESPONSE_PARSERS = ("auto", "qwen3", "minimax_m2", "glm47") SEMANTIC_ENGINE_ARGS = { "sparse_method", "deltakv_checkpoint_path", diff --git a/src/sparsevllm/entrypoints/openai/serving/response_parsing.py b/src/sparsevllm/entrypoints/openai/serving/response_parsing.py index 9849c74a..a555bfa7 100644 --- a/src/sparsevllm/entrypoints/openai/serving/response_parsing.py +++ b/src/sparsevllm/entrypoints/openai/serving/response_parsing.py @@ -46,6 +46,52 @@ "optional": True, } +_GLM_RESPONSE_TEMPLATE = { + "version": 1, + "defaults": {"role": "assistant"}, + "start_anchor": "<|assistant|>", + "fields": { + "thinking": { + # GLM writes for thinking mode and a leading + # when thinking is disabled. The zero-width alternative consumes + # that template-provided close without exposing it as content. + "open_pattern": r"\s*|(?=)", + "close": "", + "content": "text", + "optional": True, + }, + "tool_calls": { + "open_pattern": ( + r"\s*(?P[^<]*?)\s*" + r"(?=|)" + ), + "close": "", + "content": "xml-inline", + "content_args": { + "tag_pattern": ( + r"\s*(?P.*?)\s*\s*" + r"\s*(?P.*?)\s*" + ), + "value_parser": { + "name": "json", + "args": {"allow_non_json": True}, + }, + }, + "transform": { + "type": "function", + "function": {"name": "{name}", "arguments": "{content}"}, + }, + "repeats": True, + "optional": True, + }, + "content": { + "close_pattern": r"<\|endoftext\|>||\Z", + "content": "text", + "optional": True, + }, + }, +} + _MINIMAX_RESPONSE_TEMPLATE_BASE = { "version": 1, "defaults": {"role": "assistant"}, @@ -152,6 +198,12 @@ def from_tokenizer(cls, tokenizer: Any) -> "TransformersResponseParser | None": if "" not in chat_template and "" not in chat_template: return None + if all( + marker in chat_template + for marker in ("<|assistant|>", "", "") + ): + return cls(tokenizer, _GLM_RESPONSE_TEMPLATE) + if "" in chat_template and "' '' ) + if glm_tools: + self.chat_template = ( + "<|assistant|>" + ) def parse_response(self, response, schema, *, prefix=None): from transformers.utils.chat_parsing import parse_response @@ -50,13 +54,16 @@ def get_response_parser(self, response_template=None, *, prefix=None): return ResponseParser(response_template, prefix=prefix) -def _transformers_response_parser(*, xml_tools=False, minimax_tools=False): +def _transformers_response_parser( + *, xml_tools=False, minimax_tools=False, glm_tools=False +): from sparsevllm.entrypoints.openai.serving.response_parsing import TransformersResponseParser parser = TransformersResponseParser.from_tokenizer( _TransformersResponseTokenizer( xml_tools=xml_tools, minimax_tools=minimax_tools, + glm_tools=glm_tools, ) ) assert parser is not None @@ -198,6 +205,12 @@ def test_response_parser_cli_replaces_reasoning_parser(self): ["--model", "/tmp/model", "--reasoning-parser", "minimax_m2"] ) + for alias in ("auto", "glm47"): + args = parser.parse_args( + ["--model", "/tmp/model", "--response-parser", alias] + ) + self.assertEqual(args.response_parser, alias) + def test_incremental_detokenizer_waits_for_complete_unicode(self): from sparsevllm.entrypoints.openai.detokenizer import IncrementalDetokenizer @@ -5470,6 +5483,94 @@ def test_transformers_decides_xml_tool_delimiter_handling(self): "arguments": '{"command":"pwd"}', }) + def test_glm_response_parser_handles_thinking_and_plain_content(self): + parser = _transformers_response_parser(glm_tools=True) + + parsed = parser.parse( + "先分析🙂最终答案<|endoftext|>", + prefix="[gMASK]<|assistant|>", + parse_tools=False, + ) + nonthinking = parser.parse( + "直接回答<|endoftext|>", + prefix="[gMASK]<|assistant|>", + parse_tools=False, + ) + + self.assertEqual(parsed.reasoning_content, "先分析🙂") + self.assertEqual(parsed.content, "最终答案") + self.assertEqual(nonthinking.reasoning_content, None) + self.assertEqual(nonthinking.content, "直接回答") + + def test_glm_response_parser_handles_parallel_tool_calls(self): + parsed = _transformers_response_parser(glm_tools=True).parse( + "分析" + "天气查询" + "城市北京" + "天数2" + "" + "echo" + "text你好🙂" + "<|endoftext|>", + prefix="[gMASK]<|assistant|>", + parse_tools=True, + ) + + self.assertEqual(parsed.reasoning_content, "分析") + self.assertEqual(parsed.content, "") + self.assertEqual( + [call["function"]["name"] for call in parsed.tool_calls], + ["天气查询", "echo"], + ) + self.assertEqual( + parsed.tool_calls[0]["function"]["arguments"], + '{"城市":"北京","天数":2}', + ) + self.assertEqual( + parsed.tool_calls[1]["function"]["arguments"], + '{"text":"你好🙂"}', + ) + + def test_glm_response_stream_parser_handles_split_tags(self): + parser = _transformers_response_parser(glm_tools=True).stream( + prefix="[gMASK]<|assistant|>", + parse_tools=True, + ) + deltas = [] + for chunk in [ + "推", + "理答", + "案天气城市北", + "京<|endoftext|>", + ]: + deltas.extend(parser.feed(chunk)) + deltas.extend(parser.finish()) + + reasoning = "".join( + delta.get("reasoning_content", "") for delta in deltas + ) + content = "".join(delta.get("content", "") for delta in deltas) + calls = [delta["tool_calls"][0] for delta in deltas if "tool_calls" in delta] + self.assertEqual(reasoning, "推理") + self.assertEqual(content, "答案") + self.assertEqual(calls[0]["function"]["name"], "天气") + self.assertEqual(calls[0]["function"]["arguments"], '{"城市":"北京"}') + + def test_glm_response_parser_rejects_empty_tool_name(self): + from sparsevllm.entrypoints.openai.serving.response_parsing import ( + ResponseParseError, + ) + + with self.assertRaisesRegex(ResponseParseError, "non-empty function name"): + _transformers_response_parser(glm_tools=True).parse( + "", + prefix="[gMASK]<|assistant|>", + parse_tools=True, + ) + def test_minimax_tool_calls_parse_reasoning_and_parallel_invokes(self): from sparsevllm.entrypoints.openai.api_server import _response_output_items From ba4399778879046acac297617073fd7416395de1 Mon Sep 17 00:00:00 2001 From: Codex Date: Tue, 4 Aug 2026 23:56:10 +0800 Subject: [PATCH 09/48] perf: cache mla prefill validation --- src/sparsevllm/layers/mla_attention.py | 86 ++++++++++++++++++++------ tests/test_mla_attention_layer.py | 52 ++++++++++++++++ 2 files changed, 119 insertions(+), 19 deletions(-) diff --git a/src/sparsevllm/layers/mla_attention.py b/src/sparsevllm/layers/mla_attention.py index 3a4b0ee0..6d7f9b0e 100644 --- a/src/sparsevllm/layers/mla_attention.py +++ b/src/sparsevllm/layers/mla_attention.py @@ -34,6 +34,7 @@ class MlaPrefillHistory: packed_slots: torch.Tensor local_req_indices: torch.Tensor context_lens: torch.Tensor + context_lengths: tuple[int, ...] max_context_len: int required_workspace_bytes: int @@ -64,6 +65,7 @@ class _MlaPrefillPlan: packed_offsets: torch.Tensor packed_slots: torch.Tensor local_req_indices: torch.Tensor + context_lengths: tuple[int, ...] total_visible_tokens: int max_context_len: int required_workspace_bytes: int @@ -84,6 +86,33 @@ def matches( ) +@dataclass(frozen=True, slots=True) +class _MlaPrefillQueryPlan: + """Step-local query-packing validation shared by every MLA layer.""" + + validation_scope: object + source_context_lens: torch.Tensor + query_tokens: int + + def matches( + self, + validation_scope: object, + context_lens: torch.Tensor, + query_tokens: int, + ) -> bool: + return ( + self.validation_scope is validation_scope + and self.source_context_lens is context_lens + and self.query_tokens == int(query_tokens) + ) + + +def _host_int_values(tensor: torch.Tensor) -> tuple[int, ...]: + """Synchronize an integer tensor once inside a validation scope.""" + + return tuple(int(value) for value in tensor.tolist()) + + def estimate_mla_prefill_workspace_bytes( *, total_visible_tokens: int, @@ -172,6 +201,7 @@ def __init__( ) self.prefill_backend = TritonAttentionBackend() self._prefill_plan: _MlaPrefillPlan | None = None + self._prefill_query_plan: _MlaPrefillQueryPlan | None = None @classmethod def bind( @@ -280,7 +310,7 @@ def _get_prefill_plan( f"batch={batch_size} max_batch_size={self.max_batch_size}." ) - lengths = [int(value) for value in meta.context_lens.tolist()] + lengths = _host_int_values(meta.context_lens) if any(length < 0 for length in lengths): raise ValueError( f"MLA prefill context lengths must be non-negative: {lengths}." @@ -353,6 +383,7 @@ def _get_prefill_plan( packed_offsets=packed_offsets, packed_slots=packed_slots, local_req_indices=local_req_indices, + context_lengths=lengths, total_visible_tokens=total_visible_tokens, max_context_len=max_context_len, required_workspace_bytes=required_bytes, @@ -406,6 +437,7 @@ def prepare_prefill_history( packed_slots=plan.packed_slots, local_req_indices=plan.local_req_indices, context_lens=meta.context_lens, + context_lengths=plan.context_lengths, max_context_len=plan.max_context_len, required_workspace_bytes=plan.required_workspace_bytes, ) @@ -513,24 +545,40 @@ def run_prefill( f"{name} must be int32 on {self.device}, got " f"{tensor.device}/{tensor.dtype}." ) - chunks = [int(value) for value in chunk_lens.tolist()] - starts = [int(value) for value in b_start_loc.tolist()] - expected_starts: list[int] = [] - cursor = 0 - for chunk in chunks: - expected_starts.append(cursor) - cursor += chunk - if starts != expected_starts or cursor != int(q.shape[0]): - raise ValueError( - "MLA prefill query packing is inconsistent: " - f"starts={starts} expected_starts={expected_starts} " - f"chunk_tokens={cursor} q_tokens={int(q.shape[0])}." - ) - contexts = [int(value) for value in history.context_lens.tolist()] - if any(chunk <= 0 or chunk > context for chunk, context in zip(chunks, contexts)): - raise ValueError( - "MLA prefill chunk lengths must be positive and no larger than " - f"their contexts: chunks={chunks} contexts={contexts}." + validation_scope = get_context().attention_validation_scope + query_tokens = int(q.shape[0]) + cached_query_plan = self._prefill_query_plan + if cached_query_plan is None or not cached_query_plan.matches( + validation_scope, + history.context_lens, + query_tokens, + ): + chunks = _host_int_values(chunk_lens) + starts = _host_int_values(b_start_loc) + expected_starts: list[int] = [] + cursor = 0 + for chunk in chunks: + expected_starts.append(cursor) + cursor += chunk + if starts != tuple(expected_starts) or cursor != query_tokens: + raise ValueError( + "MLA prefill query packing is inconsistent: " + f"starts={starts} expected_starts={expected_starts} " + f"chunk_tokens={cursor} q_tokens={query_tokens}." + ) + contexts = history.context_lengths + if any( + chunk <= 0 or chunk > context + for chunk, context in zip(chunks, contexts) + ): + raise ValueError( + "MLA prefill chunk lengths must be positive and no larger " + f"than their contexts: chunks={chunks} contexts={contexts}." + ) + self._prefill_query_plan = _MlaPrefillQueryPlan( + validation_scope=validation_scope, + source_context_lens=history.context_lens, + query_tokens=query_tokens, ) explicit_view = PrefillComputeView( diff --git a/tests/test_mla_attention_layer.py b/tests/test_mla_attention_layer.py index 8b953814..f4692e39 100644 --- a/tests/test_mla_attention_layer.py +++ b/tests/test_mla_attention_layer.py @@ -5,6 +5,7 @@ import pytest import torch +import sparsevllm.layers.mla_attention as mla_attention_module from sparsevllm.engine.cache_manager import ( AttentionViewMeta, ExplicitKVPayload, @@ -259,6 +260,57 @@ def test_mla_prefill_reuses_validated_packing_across_layers() -> None: assert first.packed_slots is second.packed_slots +def test_mla_prefill_reuses_query_validation_across_layers() -> None: + attention = _attention() + view = _view( + torch.empty(2, 1, 512, dtype=torch.bfloat16), + torch.empty(2, 1, 64, dtype=torch.bfloat16), + torch.tensor([[0, 1]], dtype=torch.int32), + torch.tensor([0], dtype=torch.int32), + torch.tensor([2], dtype=torch.int32), + ) + with patch("sparsevllm.layers.mla_attention.gather_latent_history"): + history = attention.prepare_prefill_history(view) + workset = _expand_history(attention, history) + q = torch.empty(2, 5, 256, dtype=torch.bfloat16) + starts = torch.tensor([0], dtype=torch.int32) + chunks = torch.tensor([2], dtype=torch.int32) + + with ( + patch( + "sparsevllm.layers.mla_attention._host_int_values", + wraps=mla_attention_module._host_int_values, + ) as host_values, + patch.object( + attention.prefill_backend, + "run_prefill", + return_value=torch.empty_like(q), + ), + ): + attention.run_prefill( + q, + workset, + b_start_loc=starts, + chunk_lens=chunks, + ) + attention.run_prefill( + q, + workset, + b_start_loc=starts, + chunk_lens=chunks, + ) + assert host_values.call_count == 2 + + reset_context() + attention.run_prefill( + q, + workset, + b_start_loc=starts, + chunk_lens=chunks, + ) + assert host_values.call_count == 4 + + CUDA_REQUIRED = pytest.mark.skipif( not torch.cuda.is_available(), reason="CUDA is required for MLA full-history prefill tests", From 1209d15a66ae989f5c6bdab5dff78f385e47f9ac Mon Sep 17 00:00:00 2001 From: Codex Date: Tue, 4 Aug 2026 23:58:33 +0800 Subject: [PATCH 10/48] refactor: inject glm mla runtime --- .../engine/cache_manager/storage/__init__.py | 16 ++- src/sparsevllm/engine/model_runner.py | 36 ++++--- src/sparsevllm/models/glm4_moe_lite.py | 99 +++++++++++++------ tests/test_glm4_moe_lite.py | 35 +++++-- 4 files changed, 134 insertions(+), 52 deletions(-) diff --git a/src/sparsevllm/engine/cache_manager/storage/__init__.py b/src/sparsevllm/engine/cache_manager/storage/__init__.py index 8c5f20d6..bb1a86ce 100644 --- a/src/sparsevllm/engine/cache_manager/storage/__init__.py +++ b/src/sparsevllm/engine/cache_manager/storage/__init__.py @@ -1,10 +1,12 @@ from __future__ import annotations -from typing import Any +from typing import TYPE_CHECKING, Any from .base import AttentionCacheStorage, CacheLayout from .explicit_kv import ExplicitKVStorage -from .mla_latent import MlaLatentStorage + +if TYPE_CHECKING: + from .mla_latent import MlaLatentStorage def create_attention_cache_storage( @@ -34,6 +36,8 @@ def create_attention_cache_storage( dtype=dtype, ) if layout is CacheLayout.MLA_LATENT: + from .mla_latent import MlaLatentStorage + return MlaLatentStorage( kv_lora_rank=int(config.hf_config.kv_lora_rank), rope_dim=int(config.hf_config.qk_rope_head_dim), @@ -42,6 +46,14 @@ def create_attention_cache_storage( raise AssertionError(f"Unhandled attention cache layout: {layout!r}") +def __getattr__(name: str): + if name == "MlaLatentStorage": + from .mla_latent import MlaLatentStorage + + return MlaLatentStorage + raise AttributeError(name) + + __all__ = [ "AttentionCacheStorage", "CacheLayout", diff --git a/src/sparsevllm/engine/model_runner.py b/src/sparsevllm/engine/model_runner.py index 0d486071..1b46485f 100644 --- a/src/sparsevllm/engine/model_runner.py +++ b/src/sparsevllm/engine/model_runner.py @@ -65,12 +65,12 @@ Qwen35MoeForCausalLM = None -def _create_model(hf_config, model_spec: ModelSpec): +def _create_model(hf_config, model_spec: ModelSpec, **kwargs): class_name = model_spec.runtime_class_name model_class = globals().get(class_name) if model_class is None: raise ImportError(f"{class_name} is unavailable for {model_spec.name}.") - return model_class(hf_config) + return model_class(hf_config, **kwargs) TP_SHM_NAME_PREFIX = "sparsevllm_" @@ -165,21 +165,29 @@ def __init__( "decode_cuda_graph", bool(getattr(config, "decode_cuda_graph", False)), ) + model_kwargs = {} if config.model_spec.runtime_class_name == "Glm4MoeLiteForCausalLM": - setattr(hf_config, "sparsevllm_device", self.device) - setattr( - hf_config, - "sparsevllm_max_batch_size", - max(config.max_num_seqs_in_batch, config.max_decoding_seqs), - ) - setattr( - hf_config, - "sparsevllm_mla_prefill_workspace_bytes", - config.mla_prefill_workspace_bytes, + from sparsevllm.models.glm4_moe_lite import ( + build_glm4_moe_lite_mla_attention, ) - setattr(hf_config, "sparsevllm_expect_mtp_weights", not config.tiny_random) + + model_kwargs = { + "mla_attention": build_glm4_moe_lite_mla_attention( + hf_config, + device=self.device, + max_batch_size=max( + config.max_num_seqs_in_batch, + config.max_decoding_seqs, + ), + prefill_workspace_bytes=config.mla_prefill_workspace_bytes, + decode_cuda_graph=config.decode_cuda_graph, + ), + "mlp_chunk_size": config.mlp_chunk_size, + "decode_cuda_graph": config.decode_cuda_graph, + "expect_mtp_weights": not config.tiny_random, + } - self.model = _create_model(hf_config, config.model_spec) + self.model = _create_model(hf_config, config.model_spec, **model_kwargs) if config.tiny_random: from sparsevllm.debug.tiny_random import initialize_sparse_model diff --git a/src/sparsevllm/models/glm4_moe_lite.py b/src/sparsevllm/models/glm4_moe_lite.py index 240f0670..9452ba2d 100644 --- a/src/sparsevllm/models/glm4_moe_lite.py +++ b/src/sparsevllm/models/glm4_moe_lite.py @@ -41,18 +41,16 @@ _MTP_PREFIX = f"model.layers.{_MTP_LAYER_INDEX}." -def _build_mla_attention(config: Glm4MoeLiteConfig) -> MLAAttention: - device = getattr(config, "sparsevllm_device", None) - max_batch_size = int(getattr(config, "sparsevllm_max_batch_size", 0) or 0) - prefill_workspace_bytes = int( - getattr(config, "sparsevllm_mla_prefill_workspace_bytes", 0) or 0 - ) - if device is None or max_batch_size <= 0 or prefill_workspace_bytes <= 0: - raise ValueError( - "GLM MLA construction requires ModelRunner runtime hints: " - "sparsevllm_device, sparsevllm_max_batch_size, and " - "sparsevllm_mla_prefill_workspace_bytes." - ) +def build_glm4_moe_lite_mla_attention( + config: Glm4MoeLiteConfig, + *, + device: torch.device | str, + max_batch_size: int, + prefill_workspace_bytes: int, + decode_cuda_graph: bool, +) -> MLAAttention: + """Bind the one process-local MLA operator from explicit runtime inputs.""" + parallel_context = get_parallel_context() activation_dtype = model_activation_dtype(config) spec = MlaAttentionOpSpec( @@ -64,7 +62,7 @@ def _build_mla_attention(config: Glm4MoeLiteConfig) -> MLAAttention: activation_dtype=activation_dtype, cache_dtype=activation_dtype, tp_size=int(parallel_context.attention_tp_size), - cuda_graph=bool(getattr(config, "decode_cuda_graph", False)), + cuda_graph=bool(decode_cuda_graph), ) return MLAAttention.bind( spec=spec, @@ -81,6 +79,8 @@ def __init__( self, config: Glm4MoeLiteConfig, mla_attention: MLAAttention, + *, + projection_chunk_size: int, ) -> None: super().__init__() self.mla_attention = mla_attention @@ -93,7 +93,7 @@ def __init__( self.qk_rope_head_dim = int(config.qk_rope_head_dim) self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim self.v_head_dim = int(config.v_head_dim) - self.proj_chunk_size = int(getattr(config, "mlp_chunk_size", 16384)) + self.proj_chunk_size = int(projection_chunk_size) if self.proj_chunk_size <= 0: raise ValueError( f"mlp_chunk_size must be positive, got {self.proj_chunk_size}." @@ -356,7 +356,12 @@ def forward( class Glm4MoeLitePackedExperts(PackedMoeExperts): - def __init__(self, config: Glm4MoeLiteConfig) -> None: + def __init__( + self, + config: Glm4MoeLiteConfig, + *, + decode_cuda_graph: bool, + ) -> None: super().__init__( num_experts=int(config.n_routed_experts), hidden_size=int(config.hidden_size), @@ -364,7 +369,7 @@ def __init__(self, config: Glm4MoeLiteConfig) -> None: top_k=int(config.num_experts_per_tok), activation_dtype=model_activation_dtype(config), fp8_enabled=False, - cuda_graph=bool(getattr(config, "decode_cuda_graph", False)), + cuda_graph=bool(decode_cuda_graph), routing_method="biased_sigmoid", model_label="GLM-4.7-Flash", provider_resolver=resolve_moe_provider, @@ -373,16 +378,25 @@ def __init__(self, config: Glm4MoeLiteConfig) -> None: class Glm4MoeLiteSparseMoeBlock(nn.Module): - def __init__(self, config: Glm4MoeLiteConfig) -> None: + def __init__( + self, + config: Glm4MoeLiteConfig, + *, + mlp_chunk_size: int, + decode_cuda_graph: bool, + ) -> None: super().__init__() self.parallel_context = get_parallel_context() - self.mlp_chunk_size = int(getattr(config, "mlp_chunk_size", 16384)) + self.mlp_chunk_size = int(mlp_chunk_size) if self.mlp_chunk_size <= 0: raise ValueError( f"mlp_chunk_size must be positive, got {self.mlp_chunk_size}." ) self.gate = Glm4MoeLiteRouter(config) - self.experts = Glm4MoeLitePackedExperts(config) + self.experts = Glm4MoeLitePackedExperts( + config, + decode_cuda_graph=decode_cuda_graph, + ) self.shared_experts = Qwen3MLP( hidden_size=int(config.hidden_size), intermediate_size=( @@ -423,9 +437,16 @@ def __init__( config: Glm4MoeLiteConfig, layer_idx: int, mla_attention: MLAAttention, + *, + mlp_chunk_size: int, + decode_cuda_graph: bool, ) -> None: super().__init__() - self.self_attn = Glm4MoeLiteAttention(config, mla_attention) + self.self_attn = Glm4MoeLiteAttention( + config, + mla_attention, + projection_chunk_size=mlp_chunk_size, + ) layer_types = list(config.mlp_layer_types) if len(layer_types) != int(config.num_hidden_layers): raise ValueError( @@ -438,11 +459,15 @@ def __init__( hidden_size=int(config.hidden_size), intermediate_size=int(config.intermediate_size), hidden_act=str(config.hidden_act), - mlp_chunk_size=int(getattr(config, "mlp_chunk_size", 16384)), + mlp_chunk_size=int(mlp_chunk_size), quantization=None, ) elif layer_type == "sparse": - self.mlp = Glm4MoeLiteSparseMoeBlock(config) + self.mlp = Glm4MoeLiteSparseMoeBlock( + config, + mlp_chunk_size=mlp_chunk_size, + decode_cuda_graph=decode_cuda_graph, + ) else: raise ValueError( f"Unsupported GLM MLP layer type at layer {layer_idx}: {layer_type!r}." @@ -478,6 +503,9 @@ def __init__( self, config: Glm4MoeLiteConfig, mla_attention: MLAAttention, + *, + mlp_chunk_size: int, + decode_cuda_graph: bool, ) -> None: super().__init__() self.config = config @@ -498,7 +526,13 @@ def __init__( ) self.layers = nn.ModuleList( [ - Glm4MoeLiteDecoderLayer(config, layer_idx, mla_attention) + Glm4MoeLiteDecoderLayer( + config, + layer_idx, + mla_attention, + mlp_chunk_size=mlp_chunk_size, + decode_cuda_graph=decode_cuda_graph, + ) for layer_idx in range(int(config.num_hidden_layers)) ] ) @@ -602,20 +636,24 @@ def __init__( self, config: Glm4MoeLiteConfig, *, - mla_attention: MLAAttention | None = None, + mla_attention: MLAAttention, + mlp_chunk_size: int, + decode_cuda_graph: bool, + expect_mtp_weights: bool, ) -> None: super().__init__() self.config = config self.parallel_context = get_parallel_context() - if mla_attention is None: - mla_attention = _build_mla_attention(config) - self.model = Glm4MoeLiteModel(config, mla_attention) + self.model = Glm4MoeLiteModel( + config, + mla_attention, + mlp_chunk_size=mlp_chunk_size, + decode_cuda_graph=decode_cuda_graph, + ) self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size) if config.tie_word_embeddings: self.lm_head.weight.data = self.model.embed_tokens.weight.data - self.expect_mtp_weights = bool( - getattr(config, "sparsevllm_expect_mtp_weights", False) - ) + self.expect_mtp_weights = bool(expect_mtp_weights) self._intentionally_skipped_expert_weights: set[str] = set() self._intentionally_skipped_mtp_weights: set[str] = set() @@ -913,4 +951,5 @@ def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: "Glm4MoeLitePackedExperts", "Glm4MoeLiteRouter", "Glm4MoeLiteSparseMoeBlock", + "build_glm4_moe_lite_mla_attention", ] diff --git a/tests/test_glm4_moe_lite.py b/tests/test_glm4_moe_lite.py index 1e913ade..9626a35d 100644 --- a/tests/test_glm4_moe_lite.py +++ b/tests/test_glm4_moe_lite.py @@ -64,7 +64,6 @@ def _config(**overrides) -> Glm4MoeLiteConfig: config.mlp_chunk_size = 8 config.quantization_config = QuantizationConfig.disabled() config.decode_cuda_graph = False - config.sparsevllm_expect_mtp_weights = False return config @@ -136,6 +135,9 @@ def _model(config=None, *, tp_rank: int = 0, tp_size: int = 1): return Glm4MoeLiteForCausalLM( config, mla_attention=_fake_mla(tp_size), + mlp_chunk_size=config.mlp_chunk_size, + decode_cuda_graph=config.decode_cuda_graph, + expect_mtp_weights=False, ) @@ -181,7 +183,11 @@ def test_glm_tp_projection_slices_follow_local_heads() -> None: config = _config() context = _tp_context(tp_rank=2, tp_size=4) with _construction_context(context): - attention = Glm4MoeLiteAttention(config, _fake_mla(tp_size=4)) + attention = Glm4MoeLiteAttention( + config, + _fake_mla(tp_size=4), + projection_chunk_size=config.mlp_chunk_size, + ) q_source = ( torch.arange(20 * 256 * 768).remainder(127).to(torch.bfloat16) @@ -206,7 +212,11 @@ def test_glm_decode_absorption_and_value_reconstruction_match_linear_algebra() - config = _config() context = _tp_context() with _construction_context(context): - attention = Glm4MoeLiteAttention(config, _fake_mla()) + attention = Glm4MoeLiteAttention( + config, + _fake_mla(), + projection_chunk_size=config.mlp_chunk_size, + ) torch.manual_seed(19) attention.kv_b_proj.weight.data.normal_(mean=0.0, std=0.02) q_nope = torch.randn(3, 20, 192, dtype=torch.bfloat16) @@ -264,7 +274,13 @@ def test_tiny_transformers_weights_load_through_strict_glm_mapping() -> None: config = _config() context = _tp_context() with _construction_context(context): - model = Glm4MoeLiteForCausalLM(config, mla_attention=_fake_mla()) + model = Glm4MoeLiteForCausalLM( + config, + mla_attention=_fake_mla(), + mlp_chunk_size=config.mlp_chunk_size, + decode_cuda_graph=config.decode_cuda_graph, + expect_mtp_weights=False, + ) initialize_sparse_model(model, config, seed=29) reference = build_tiny_random_hf_model(config, seed=29) reference_experts = reference.model.layers[1].mlp.experts @@ -292,8 +308,15 @@ def test_tiny_transformers_weights_load_through_strict_glm_mapping() -> None: def test_glm_mtp_skip_set_is_exact() -> None: config = _config() - config.sparsevllm_expect_mtp_weights = True - model = _model(config) + context = _tp_context() + with _construction_context(context): + model = Glm4MoeLiteForCausalLM( + config, + mla_attention=_fake_mla(), + mlp_chunk_size=config.mlp_chunk_size, + decode_cuda_graph=config.decode_cuda_graph, + expect_mtp_weights=True, + ) dense_names = { name for name, _ in model.named_parameters() From 0fbb579cf933f91efdcebdf2a6c3c66d0b51f722 Mon Sep 17 00:00:00 2001 From: Codex Date: Wed, 5 Aug 2026 00:02:40 +0800 Subject: [PATCH 11/48] fix: correct mla memory accounting --- src/sparsevllm/engine/cache_manager/base.py | 33 ++++- src/sparsevllm/engine/model_runner.py | 1 + src/sparsevllm/layers/mla_attention.py | 141 ++++++++++++++++++-- src/sparsevllm/models/glm4_moe_lite.py | 19 ++- tests/test_attention_cache_storage.py | 23 +++- tests/test_glm4_moe_lite.py | 2 + tests/test_mla_attention_layer.py | 43 ++++-- 7 files changed, 233 insertions(+), 29 deletions(-) diff --git a/src/sparsevllm/engine/cache_manager/base.py b/src/sparsevllm/engine/cache_manager/base.py index a698424e..54e92b3f 100644 --- a/src/sparsevllm/engine/cache_manager/base.py +++ b/src/sparsevllm/engine/cache_manager/base.py @@ -1407,7 +1407,38 @@ def _dense_baseline_slots(self) -> int: def _dense_baseline_bytes(self) -> int: dtype_size = self._cache_slot_dtype_size() slots = self._dense_baseline_slots() - return int(slots * self.num_kv_layers * 2 * self.num_kv_heads * self.head_dim * dtype_size) + storage = getattr(self, "attention_cache_storage", None) + layout = getattr(storage, "layout", None) + layout_name = str(getattr(layout, "value", layout) or "") + if layout_name == "mla_latent": + hf_config = self.hf_config + global_heads = int(hf_config.num_attention_heads) + attention_tp_size = int( + getattr( + getattr(self, "parallel_context", None), + "attention_tp_size", + getattr(self, "tp_size", 1), + ) + ) + if global_heads % attention_tp_size != 0: + raise ValueError( + "MLA dense baseline requires attention heads divisible by TP: " + f"heads={global_heads} tp={attention_tp_size}." + ) + local_heads = global_heads // attention_tp_size + qk_head_dim = int(hf_config.qk_nope_head_dim) + int( + hf_config.qk_rope_head_dim + ) + value_head_dim = int(hf_config.v_head_dim) + values_per_token = local_heads * (qk_head_dim + value_head_dim) + else: + values_per_token = 2 * self.num_kv_heads * self.head_dim + return int( + slots + * self.num_kv_layers + * values_per_token + * dtype_size + ) @staticmethod def _tensor_storage_key(tensor: torch.Tensor) -> tuple[Any, ...]: diff --git a/src/sparsevllm/engine/model_runner.py b/src/sparsevllm/engine/model_runner.py index 1b46485f..78b606f1 100644 --- a/src/sparsevllm/engine/model_runner.py +++ b/src/sparsevllm/engine/model_runner.py @@ -181,6 +181,7 @@ def __init__( ), prefill_workspace_bytes=config.mla_prefill_workspace_bytes, decode_cuda_graph=config.decode_cuda_graph, + projection_chunk_size=config.mlp_chunk_size, ), "mlp_chunk_size": config.mlp_chunk_size, "decode_cuda_graph": config.decode_cuda_graph, diff --git a/src/sparsevllm/layers/mla_attention.py b/src/sparsevllm/layers/mla_attention.py index 6d7f9b0e..53f37765 100644 --- a/src/sparsevllm/layers/mla_attention.py +++ b/src/sparsevllm/layers/mla_attention.py @@ -61,6 +61,7 @@ class _MlaPrefillPlan: source_req_indices: torch.Tensor source_context_lens: torch.Tensor source_max_context_len: int | None + source_query_tokens: int cache_slot_count: int packed_offsets: torch.Tensor packed_slots: torch.Tensor @@ -75,6 +76,7 @@ def matches( validation_scope: object, meta: AttentionViewMeta, cache_slot_count: int, + query_tokens: int, ) -> bool: return ( self.validation_scope is validation_scope @@ -82,6 +84,7 @@ def matches( and self.source_req_indices is meta.req_indices and self.source_context_lens is meta.context_lens and self.source_max_context_len == meta.max_context_len + and self.source_query_tokens == int(query_tokens) and self.cache_slot_count == int(cache_slot_count) ) @@ -116,49 +119,133 @@ def _host_int_values(tensor: torch.Tensor) -> tuple[int, ...]: def estimate_mla_prefill_workspace_bytes( *, total_visible_tokens: int, + query_tokens: int, batch_size: int, max_context_len: int, local_q_heads: int, + kv_lora_rank: int, + rope_dim: int, + qk_head_dim: int, + value_head_dim: int, + hidden_size: int, + projection_chunk_size: int, activation_dtype: torch.dtype, cache_dtype: torch.dtype, ) -> int: - """Conservatively bound full-history gather and K/V expansion storage.""" + """Bound the peak modeled transient storage for MLA full-history prefill.""" values = { "total_visible_tokens": total_visible_tokens, + "query_tokens": query_tokens, "batch_size": batch_size, "max_context_len": max_context_len, "local_q_heads": local_q_heads, + "kv_lora_rank": kv_lora_rank, + "rope_dim": rope_dim, + "qk_head_dim": qk_head_dim, + "value_head_dim": value_head_dim, + "hidden_size": hidden_size, + "projection_chunk_size": projection_chunk_size, } for name, value in values.items(): if int(value) < 0: raise ValueError(f"{name} must be non-negative, got {value}.") - if batch_size == 0 or local_q_heads == 0: - raise ValueError("batch_size and local_q_heads must be positive.") + positive_values = ( + "total_visible_tokens", + "query_tokens", + "batch_size", + "local_q_heads", + "kv_lora_rank", + "rope_dim", + "qk_head_dim", + "value_head_dim", + "hidden_size", + "projection_chunk_size", + ) + for name in positive_values: + if int(values[name]) == 0: + raise ValueError(f"{name} must be positive, got 0.") + if int(query_tokens) > int(total_visible_tokens): + raise ValueError( + "query_tokens cannot exceed total_visible_tokens, got " + f"{query_tokens} > {total_visible_tokens}." + ) + qk_nope_head_dim = int(qk_head_dim) - int(rope_dim) + if qk_nope_head_dim <= 0: + raise ValueError( + "qk_head_dim must be larger than rope_dim, got " + f"{qk_head_dim} and {rope_dim}." + ) cache_element_size = torch.empty((), dtype=cache_dtype).element_size() activation_element_size = torch.empty( (), dtype=activation_dtype, ).element_size() - gathered_values = int(total_visible_tokens) * (512 + 64) - # kv_b projection materializes 192 K-noPE + 256 V values per head. The - # final K concatenation additionally materializes 192+64 values per head; - # expanded_v may remain a view of the projection output. - expanded_values = ( - int(total_visible_tokens) - * int(local_q_heads) - * ((192 + 256) + 256) + visible_tokens = int(total_visible_tokens) + current_tokens = int(query_tokens) + heads = int(local_q_heads) + projected_width = qk_nope_head_dim + int(value_head_dim) + gathered_bytes = ( + visible_tokens + * (int(kv_lora_rank) + int(rope_dim)) + * cache_element_size + ) + projected_bytes = ( + visible_tokens * heads * projected_width * activation_element_size + ) + projection_scratch_bytes = 0 + if visible_tokens > int(projection_chunk_size): + projection_scratch_bytes = ( + min(visible_tokens, int(projection_chunk_size)) + * heads + * projected_width + * activation_element_size + ) + expanded_k_bytes = ( + visible_tokens + * heads + * int(qk_head_dim) + * activation_element_size + ) + attention_output_bytes = ( + current_tokens + * heads + * int(value_head_dim) + * activation_element_size + ) + output_projection_scratch_bytes = ( + min(current_tokens, int(projection_chunk_size)) + * int(hidden_size) + * activation_element_size + ) + kv_projection_phase_bytes = ( + gathered_bytes + projected_bytes + projection_scratch_bytes + ) + attention_phase_bytes = ( + gathered_bytes + + projected_bytes + + expanded_k_bytes + + attention_output_bytes + ) + output_projection_phase_bytes = ( + attention_output_bytes + output_projection_scratch_bytes ) metadata_values = ( int(batch_size) * int(max_context_len) + 2 * int(batch_size) + int(max_context_len) ) + metadata_bytes = ( + metadata_values * torch.empty((), dtype=torch.int32).element_size() + ) return int( - gathered_values * cache_element_size - + expanded_values * activation_element_size - + metadata_values * torch.empty((), dtype=torch.int32).element_size() + max( + kv_projection_phase_bytes, + attention_phase_bytes, + output_projection_phase_bytes, + ) + + metadata_bytes ) @@ -176,6 +263,8 @@ def __init__( spec: MlaAttentionOpSpec, provider: MlaAttentionProvider, prefill_workspace_bytes: int, + hidden_size: int, + projection_chunk_size: int, ) -> None: self.spec = spec self.provider = provider @@ -194,6 +283,13 @@ def __init__( "MLA prefill_workspace_bytes must be positive, got " f"{self.prefill_workspace_bytes}." ) + self.hidden_size = int(hidden_size) + self.projection_chunk_size = int(projection_chunk_size) + if self.hidden_size <= 0 or self.projection_chunk_size <= 0: + raise ValueError( + "MLA hidden_size and projection_chunk_size must be positive, " + f"got {self.hidden_size} and {self.projection_chunk_size}." + ) if self.spec.qk_head_dim != self.spec.value_head_dim: raise ValueError( "The existing prefill backend requires equal QK/value widths, " @@ -211,6 +307,8 @@ def bind( device: torch.device | str, max_batch_size: int, prefill_workspace_bytes: int, + hidden_size: int, + projection_chunk_size: int, ) -> "MLAAttention": provider = resolve_mla_attention_provider( spec, @@ -221,6 +319,8 @@ def bind( spec=spec, provider=provider, prefill_workspace_bytes=prefill_workspace_bytes, + hidden_size=hidden_size, + projection_chunk_size=projection_chunk_size, ) @property @@ -265,6 +365,7 @@ def _get_prefill_plan( meta: AttentionViewMeta, *, cache_slot_count: int, + query_tokens: int, ) -> _MlaPrefillPlan: cached = self._prefill_plan validation_scope = get_context().attention_validation_scope @@ -272,6 +373,7 @@ def _get_prefill_plan( validation_scope, meta, cache_slot_count, + query_tokens, ): return cached @@ -329,9 +431,16 @@ def _get_prefill_plan( ) required_bytes = estimate_mla_prefill_workspace_bytes( total_visible_tokens=total_visible_tokens, + query_tokens=query_tokens, batch_size=batch_size, max_context_len=max_context_len, local_q_heads=self.spec.local_q_heads, + kv_lora_rank=self.spec.kv_lora_rank, + rope_dim=self.spec.rope_dim, + qk_head_dim=self.spec.qk_head_dim, + value_head_dim=self.spec.value_head_dim, + hidden_size=self.hidden_size, + projection_chunk_size=self.projection_chunk_size, activation_dtype=self.spec.activation_dtype, cache_dtype=self.spec.cache_dtype, ) @@ -379,6 +488,7 @@ def _get_prefill_plan( source_req_indices=meta.req_indices, source_context_lens=meta.context_lens, source_max_context_len=meta.max_context_len, + source_query_tokens=int(query_tokens), cache_slot_count=int(cache_slot_count), packed_offsets=packed_offsets, packed_slots=packed_slots, @@ -394,6 +504,8 @@ def _get_prefill_plan( def prepare_prefill_history( self, view: PrefillComputeView, + *, + query_tokens: int, ) -> MlaPrefillHistory: if not isinstance(view, PrefillComputeView): raise TypeError( @@ -405,6 +517,7 @@ def prepare_prefill_history( plan = self._get_prefill_plan( meta, cache_slot_count=int(payload.latent_cache.shape[0]), + query_tokens=query_tokens, ) gathered_latent = torch.empty( plan.total_visible_tokens, diff --git a/src/sparsevllm/models/glm4_moe_lite.py b/src/sparsevllm/models/glm4_moe_lite.py index 9452ba2d..1452af1e 100644 --- a/src/sparsevllm/models/glm4_moe_lite.py +++ b/src/sparsevllm/models/glm4_moe_lite.py @@ -48,6 +48,7 @@ def build_glm4_moe_lite_mla_attention( max_batch_size: int, prefill_workspace_bytes: int, decode_cuda_graph: bool, + projection_chunk_size: int, ) -> MLAAttention: """Bind the one process-local MLA operator from explicit runtime inputs.""" @@ -69,6 +70,8 @@ def build_glm4_moe_lite_mla_attention( device=device, max_batch_size=max_batch_size, prefill_workspace_bytes=prefill_workspace_bytes, + hidden_size=int(config.hidden_size), + projection_chunk_size=projection_chunk_size, ) @@ -98,6 +101,17 @@ def __init__( raise ValueError( f"mlp_chunk_size must be positive, got {self.proj_chunk_size}." ) + if self.proj_chunk_size != int(mla_attention.projection_chunk_size): + raise ValueError( + "GLM projection chunk size must match the MLA workspace bound: " + f"model={self.proj_chunk_size} " + f"mla={mla_attention.projection_chunk_size}." + ) + if int(config.hidden_size) != int(mla_attention.hidden_size): + raise ValueError( + "GLM hidden size must match the MLA workspace bound: " + f"model={config.hidden_size} mla={mla_attention.hidden_size}." + ) quantization = getattr(config, "quantization_config", None) if bool(getattr(quantization, "enabled", False)): raise NotImplementedError("GLM MLA projections do not support quantization.") @@ -216,7 +230,10 @@ def _run_attention( temp_slots = view.meta.temp_slots if context.cu_seqlens_q is None or context.cu_seqlens_q.numel() <= 1: return torch.empty_like(q) - history = self.mla_attention.prepare_prefill_history(view) + history = self.mla_attention.prepare_prefill_history( + view, + query_tokens=int(q.shape[0]), + ) expanded = self._project_kv_history(history.gathered_latent).view( history.visible_tokens, self.local_heads, diff --git a/tests/test_attention_cache_storage.py b/tests/test_attention_cache_storage.py index 64fafcde..1099e2ad 100644 --- a/tests/test_attention_cache_storage.py +++ b/tests/test_attention_cache_storage.py @@ -366,7 +366,14 @@ def test_standard_manager_delegates_payload_store_and_compute_view(): assert actual_lens is context_lens -def test_standard_manager_accounts_storage_tensors_explicitly(): +@pytest.mark.parametrize( + ("attention_tp_size", "local_heads"), + [(1, 20), (4, 5)], +) +def test_standard_manager_accounts_storage_tensors_explicitly( + attention_tp_size, + local_heads, +): storage = MlaLatentStorage( kv_lora_rank=512, rope_dim=64, @@ -381,7 +388,16 @@ def test_standard_manager_accounts_storage_tensors_explicitly(): max_num_seqs_in_gpu=1, memory_expected_savings=None, ) - manager.hf_config = SimpleNamespace(torch_dtype=torch.bfloat16) + manager.hf_config = SimpleNamespace( + torch_dtype=torch.bfloat16, + num_attention_heads=20, + qk_nope_head_dim=192, + qk_rope_head_dim=64, + v_head_dim=256, + ) + manager.parallel_context = SimpleNamespace( + attention_tp_size=attention_tp_size, + ) manager.num_layers = 2 manager.num_kv_layers = 2 manager.num_kv_heads = 4 @@ -392,6 +408,9 @@ def test_standard_manager_accounts_storage_tensors_explicitly(): assert accounting["kv_or_latent_tensor_bytes"] == 2 * 3 * 576 * 2 assert accounting["logical_live_kv_bytes"] == 2 * 2 * 576 * 2 + assert accounting["dense_baseline_bytes"] == ( + 3 * 2 * local_heads * (256 + 256) * 2 + ) assert accounting["tensor_count"] == 2 assert {item["path"] for item in accounting["tensors"]} == { "attention_cache_storage.mla_latent.0_cache", diff --git a/tests/test_glm4_moe_lite.py b/tests/test_glm4_moe_lite.py index 9626a35d..852cc01c 100644 --- a/tests/test_glm4_moe_lite.py +++ b/tests/test_glm4_moe_lite.py @@ -92,6 +92,8 @@ def _fake_mla(tp_size: int = 1): return SimpleNamespace( spec=spec, provider=SimpleNamespace(name="test"), + hidden_size=64, + projection_chunk_size=8, ) diff --git a/tests/test_mla_attention_layer.py b/tests/test_mla_attention_layer.py index f4692e39..779245b7 100644 --- a/tests/test_mla_attention_layer.py +++ b/tests/test_mla_attention_layer.py @@ -70,6 +70,8 @@ def _attention( max_batch_size=max_batch_size, ), prefill_workspace_bytes=budget, + hidden_size=64, + projection_chunk_size=8, ) @@ -145,17 +147,32 @@ def _torch_prefill( def test_mla_prefill_workspace_estimate_accounts_for_full_history() -> None: actual = estimate_mla_prefill_workspace_bytes( total_visible_tokens=11, + query_tokens=5, batch_size=2, max_context_len=7, local_q_heads=5, + kv_lora_rank=512, + rope_dim=64, + qk_head_dim=256, + value_head_dim=256, + hidden_size=64, + projection_chunk_size=4, activation_dtype=torch.bfloat16, cache_dtype=torch.bfloat16, ) gathered = 11 * (512 + 64) * 2 - expanded = 11 * 5 * ((192 + 256) + 256) * 2 + projected = 11 * 5 * (192 + 256) * 2 + projection_scratch = 4 * 5 * (192 + 256) * 2 + expanded_k = 11 * 5 * 256 * 2 + attention_output = 5 * 5 * 256 * 2 + output_projection_scratch = 4 * 64 * 2 metadata = (2 * 7 + 2 * 2 + 7) * 4 - assert actual == gathered + expanded + metadata + assert actual == max( + gathered + projected + projection_scratch, + gathered + projected + expanded_k + attention_output, + attention_output + output_projection_scratch, + ) + metadata def test_mla_prefill_rejects_wrong_payload_before_gather() -> None: @@ -176,7 +193,7 @@ def test_mla_prefill_rejects_wrong_payload_before_gather() -> None: patch("sparsevllm.layers.mla_attention.gather_latent_history") as gather, pytest.raises(TypeError, match="MlaLatentPayload"), ): - attention.prepare_prefill_history(view) + attention.prepare_prefill_history(view, query_tokens=1) gather.assert_not_called() @@ -194,7 +211,7 @@ def test_mla_prefill_budget_fails_before_allocation_or_gather() -> None: patch("sparsevllm.layers.mla_attention.gather_latent_history") as gather, pytest.raises(MemoryError, match="exceeds its configured budget"), ): - attention.prepare_prefill_history(view) + attention.prepare_prefill_history(view, query_tokens=2) gather.assert_not_called() @@ -211,6 +228,8 @@ def test_mla_attention_bind_resolves_provider_once() -> None: device="cpu", max_batch_size=8, prefill_workspace_bytes=1024, + hidden_size=64, + projection_chunk_size=8, ) assert attention.provider is provider @@ -249,10 +268,10 @@ def test_mla_prefill_reuses_validated_packing_across_layers() -> None: "sparsevllm.layers.mla_attention.gather_latent_history" ) as gather, ): - first = attention.prepare_prefill_history(first_view) - second = attention.prepare_prefill_history(second_view) + first = attention.prepare_prefill_history(first_view, query_tokens=2) + second = attention.prepare_prefill_history(second_view, query_tokens=2) reset_context() - attention.prepare_prefill_history(second_view) + attention.prepare_prefill_history(second_view, query_tokens=2) assert validate.call_count == 2 assert gather.call_count == 3 @@ -270,7 +289,7 @@ def test_mla_prefill_reuses_query_validation_across_layers() -> None: torch.tensor([2], dtype=torch.int32), ) with patch("sparsevllm.layers.mla_attention.gather_latent_history"): - history = attention.prepare_prefill_history(view) + history = attention.prepare_prefill_history(view, query_tokens=2) workset = _expand_history(attention, history) q = torch.empty(2, 5, 256, dtype=torch.bfloat16) starts = torch.tensor([0], dtype=torch.int32) @@ -337,7 +356,7 @@ def test_mla_prefill_matches_ragged_full_history_oracle() -> None: torch.tensor([2, 0], dtype=torch.int32, device="cuda"), context_lens, ) - history = attention.prepare_prefill_history(view) + history = attention.prepare_prefill_history(view, query_tokens=5) workset = _expand_history(attention, history) chunk_lens = torch.tensor([2, 3], dtype=torch.int32, device="cuda") b_start_loc = torch.tensor([0, 2], dtype=torch.int32, device="cuda") @@ -382,7 +401,8 @@ def test_mla_prefill_is_invariant_to_chunk_boundary() -> None: active_slots, request_indices, torch.tensor([6], dtype=torch.int32, device="cuda"), - ) + ), + query_tokens=6, ) full_workset = _expand_history(attention, full_history) full_output = attention.run_prefill( @@ -399,7 +419,8 @@ def test_mla_prefill_is_invariant_to_chunk_boundary() -> None: active_slots, request_indices, torch.tensor([4], dtype=torch.int32, device="cuda"), - ) + ), + query_tokens=4, ) first_output = attention.run_prefill( q[:4], From e51f88fb902f14e4325ab7104791d69550ffcfbb Mon Sep 17 00:00:00 2001 From: Codex Date: Wed, 5 Aug 2026 00:24:16 +0800 Subject: [PATCH 12/48] fix: strip terminal eos from parser --- .../entrypoints/openai/dispatcher.py | 104 ++++++++++++- tests/test_openai_api_server.py | 142 ++++++++++++++++++ 2 files changed, 240 insertions(+), 6 deletions(-) diff --git a/src/sparsevllm/entrypoints/openai/dispatcher.py b/src/sparsevllm/entrypoints/openai/dispatcher.py index 7bad90b2..351c651c 100644 --- a/src/sparsevllm/entrypoints/openai/dispatcher.py +++ b/src/sparsevllm/entrypoints/openai/dispatcher.py @@ -58,6 +58,8 @@ class _ActiveRequest: completion_token_logprobs: list[float | None] completion_top_logprobs: list[dict[int, float] | None] detokenizer: IncrementalDetokenizer + eos_token_ids: frozenset[int] = field(default_factory=frozenset) + ignore_eos: bool = False terminal: threading.Event = field(default_factory=threading.Event) emitted_text_len: int = 0 pending_token_ids: list[int] = field(default_factory=list) @@ -582,6 +584,23 @@ def _admit(self, item: _QueuedRequest, active: dict[int, _ActiveRequest]): if isinstance(item.prompt, list) else self.engine.tokenizer.encode(item.prompt) ) + requested_eos = getattr( + item.sampling_params, "eos_token_ids", () + ) or () + if isinstance(requested_eos, int): + requested_eos = (requested_eos,) + eos_token_ids = frozenset( + int(token_id) for token_id in requested_eos + ) + if not eos_token_ids: + configured_eos = getattr( + getattr(self.engine, "config", None), + "eos_token_ids", + (), + ) or () + eos_token_ids = frozenset( + int(token_id) for token_id in configured_eos + ) active[seq_id] = _ActiveRequest( index=item.index, loop=item.loop, @@ -593,6 +612,10 @@ def _admit(self, item: _QueuedRequest, active: dict[int, _ActiveRequest]): completion_token_logprobs=[], completion_top_logprobs=[], detokenizer=detokenizer, + eos_token_ids=eos_token_ids, + ignore_eos=bool( + getattr(item.sampling_params, "ignore_eos", False) + ), terminal=item.handle.terminal, chain_id=item.handle.chain_id, chain_status=item.handle.chain_status, @@ -647,6 +670,40 @@ def _admit(self, item: _QueuedRequest, active: dict[int, _ActiveRequest]): ) self._resolve_admission(item, exc) + @staticmethod + def _response_parser_raw_text( + request: _ActiveRequest, + token_ids: list[int], + raw_text: str, + ) -> str: + if ( + request.ignore_eos + or not request.eos_token_ids + or not token_ids + or int(token_ids[-1]) not in request.eos_token_ids + ): + return raw_text + + content_end = len(token_ids) + while ( + content_end > 0 + and int(token_ids[content_end - 1]) in request.eos_token_ids + ): + content_end -= 1 + return request.detokenizer.tokenizer.decode( + token_ids[:content_end], + skip_special_tokens=False, + ) + + @staticmethod + def _response_parser_delta(previous: str, current: str) -> str: + if not current.startswith(previous): + raise RuntimeError( + "Response-parser text is not append-only: " + f"previous={previous!r} current={current!r}." + ) + return current[len(previous):] + def _publish_token_deltas(self, active: dict[int, _ActiveRequest]): logprob_outputs = { seq_id: (token_logprobs, top_logprobs) @@ -664,14 +721,27 @@ def _publish_token_deltas(self, active: dict[int, _ActiveRequest]): seq_id, ([None] * len(token_ids), [None] * len(token_ids)), ) + previous_parser_raw_text = self._response_parser_raw_text( + request, + request.completion_token_ids, + request.detokenizer.raw_text, + ) request.completion_token_ids.extend(token_ids) request.completion_token_logprobs.extend(token_logprobs) request.completion_top_logprobs.extend(top_logprobs) request.pending_token_ids.extend(token_ids) request.pending_token_logprobs.extend(token_logprobs) request.pending_top_logprobs.extend(top_logprobs) - decoded = request.detokenizer.push(token_ids) - raw_text_delta = decoded.raw_text + request.detokenizer.push(token_ids) + parser_raw_text = self._response_parser_raw_text( + request, + request.completion_token_ids, + request.detokenizer.raw_text, + ) + raw_text_delta = self._response_parser_delta( + previous_parser_raw_text, + parser_raw_text, + ) full_text = request.detokenizer.text stop_index = _find_stop_index(full_text, request.stop) visible_text = full_text if stop_index is None else full_text[:stop_index] @@ -704,7 +774,11 @@ def _publish_token_deltas(self, active: dict[int, _ActiveRequest]): "type": "final", "index": request.index, "text": visible_text, - "raw_text": final.raw_text, + "raw_text": self._response_parser_raw_text( + request, + request.completion_token_ids, + final.raw_text, + ), "text_delta": visible_text[request.emitted_text_len:], "finish_reason": "stop", "prompt_tokens": ( @@ -969,7 +1043,21 @@ def _publish_finished( if request is None: continue observed = len(request.completion_token_ids) + previous_parser_raw_text = self._response_parser_raw_text( + request, + request.completion_token_ids, + request.detokenizer.raw_text, + ) final = request.detokenizer.finish(completion_token_ids) + parser_raw_text = self._response_parser_raw_text( + request, + completion_token_ids, + final.raw_text, + ) + parser_raw_text_delta = self._response_parser_delta( + previous_parser_raw_text, + parser_raw_text, + ) request.pending_token_ids.extend(completion_token_ids[observed:]) request.pending_token_logprobs.extend(token_logprobs[observed:]) request.pending_top_logprobs.extend(top_logprobs[observed:]) @@ -997,9 +1085,13 @@ def _publish_finished( value is not None for value in request.pending_token_logprobs ) or any(value is not None for value in request.pending_top_logprobs) if request.pending_token_ids and ( - text_delta or final.raw_text_delta or has_pending_logprobs + text_delta or parser_raw_text_delta or has_pending_logprobs ): - self._publish_pending_token_event(request, text_delta, final.raw_text_delta) + self._publish_pending_token_event( + request, + text_delta, + parser_raw_text_delta, + ) request.emitted_text_len = len(text) self._mark_request_terminal(seq_id, request) self._put( @@ -1008,7 +1100,7 @@ def _publish_finished( "type": "final", "index": request.index, "text": text, - "raw_text": final.raw_text, + "raw_text": parser_raw_text, "text_delta": text[request.emitted_text_len:], "finish_reason": finish_reason, "prompt_tokens": ( diff --git a/tests/test_openai_api_server.py b/tests/test_openai_api_server.py index 091b31d2..e9b6cca8 100644 --- a/tests/test_openai_api_server.py +++ b/tests/test_openai_api_server.py @@ -2618,6 +2618,148 @@ def exit(self): self.assertEqual(item["raw_text_delta"], "b") self.assertEqual(active[7].completion_token_ids, token_ids) + async def test_dispatcher_strips_terminal_eos_from_parser_text(self): + from sparsevllm.entrypoints.openai.api_server import AsyncEngineDispatcher, _ActiveRequest + from sparsevllm.entrypoints.openai.detokenizer import IncrementalDetokenizer + + tokenizer = _byte_level_tokenizer(special_tokens=[""]) + content_token_ids = tokenizer.encode("Paris") + eos_token_id = tokenizer._tokenizer.token_to_id("") + completion_token_ids = content_token_ids + [eos_token_id] + + class Engine: + def __init__(self): + self.tokenizer = tokenizer + self.last_step_token_outputs = [] + self.last_step_logprob_outputs = [] + + def exit(self): + pass + + engine = Engine() + dispatcher = AsyncEngineDispatcher(engine) + output_queue = asyncio.Queue() + active = { + 7: _ActiveRequest( + index=0, + loop=asyncio.get_running_loop(), + output_queue=output_queue, + prompt_token_ids=[10], + max_tokens=len(completion_token_ids) + 1, + stop=[], + completion_token_ids=[], + completion_token_logprobs=[], + completion_top_logprobs=[], + detokenizer=IncrementalDetokenizer(tokenizer), + eos_token_ids=frozenset({eos_token_id}), + ) + } + try: + engine.last_step_token_outputs = [(7, content_token_ids)] + engine.last_step_logprob_outputs = [ + ( + 7, + [None] * len(content_token_ids), + [None] * len(content_token_ids), + ) + ] + dispatcher._publish_token_deltas(active) + token_item = await asyncio.wait_for(output_queue.get(), timeout=1) + + engine.last_step_token_outputs = [(7, [eos_token_id])] + engine.last_step_logprob_outputs = [(7, [None], [None])] + dispatcher._publish_token_deltas(active) + self.assertTrue(output_queue.empty()) + + dispatcher._publish_finished( + active, + [ + ( + 7, + completion_token_ids, + [None] * len(completion_token_ids), + [None] * len(completion_token_ids), + ) + ], + ) + final_item = await asyncio.wait_for(output_queue.get(), timeout=1) + finally: + dispatcher.close() + + self.assertEqual(token_item["raw_text_delta"], "Paris") + self.assertEqual(final_item["raw_text"], "Paris") + self.assertEqual(final_item["token_ids"], completion_token_ids) + self.assertEqual( + final_item["completion_tokens"], + len(completion_token_ids), + ) + + async def test_dispatcher_preserves_eos_when_ignored(self): + from sparsevllm.entrypoints.openai.api_server import AsyncEngineDispatcher, _ActiveRequest + from sparsevllm.entrypoints.openai.detokenizer import IncrementalDetokenizer + + tokenizer = _byte_level_tokenizer(special_tokens=[""]) + content_token_ids = tokenizer.encode("Paris") + eos_token_id = tokenizer._tokenizer.token_to_id("") + completion_token_ids = content_token_ids + [eos_token_id] + + class Engine: + def __init__(self): + self.tokenizer = tokenizer + self.last_step_token_outputs = [ + (7, completion_token_ids) + ] + self.last_step_logprob_outputs = [ + ( + 7, + [None] * len(completion_token_ids), + [None] * len(completion_token_ids), + ) + ] + + def exit(self): + pass + + engine = Engine() + dispatcher = AsyncEngineDispatcher(engine) + output_queue = asyncio.Queue() + active = { + 7: _ActiveRequest( + index=0, + loop=asyncio.get_running_loop(), + output_queue=output_queue, + prompt_token_ids=[10], + max_tokens=len(completion_token_ids), + stop=[], + completion_token_ids=[], + completion_token_logprobs=[], + completion_top_logprobs=[], + detokenizer=IncrementalDetokenizer(tokenizer), + eos_token_ids=frozenset({eos_token_id}), + ignore_eos=True, + ) + } + try: + dispatcher._publish_token_deltas(active) + token_item = await asyncio.wait_for(output_queue.get(), timeout=1) + dispatcher._publish_finished( + active, + [ + ( + 7, + completion_token_ids, + [None] * len(completion_token_ids), + [None] * len(completion_token_ids), + ) + ], + ) + final_item = await asyncio.wait_for(output_queue.get(), timeout=1) + finally: + dispatcher.close() + + self.assertEqual(token_item["raw_text_delta"], "Paris") + self.assertEqual(final_item["raw_text"], "Paris") + async def test_dispatcher_streams_complete_unicode_with_pending_logprobs(self): from sparsevllm.entrypoints.openai.api_server import AsyncEngineDispatcher, _ActiveRequest from sparsevllm.entrypoints.openai.detokenizer import IncrementalDetokenizer From f5e885a972efbb93ba7b17e966b9d78ee895f600 Mon Sep 17 00:00:00 2001 From: Codex Date: Wed, 5 Aug 2026 00:55:41 +0800 Subject: [PATCH 13/48] fix: align parser stop boundaries --- src/sparsevllm/engine/scheduler.py | 21 ++- .../entrypoints/openai/detokenizer.py | 99 +++++++++- .../entrypoints/openai/dispatcher.py | 171 +++++++++++------- src/sparsevllm/sampling_params.py | 31 ++++ tests/test_openai_api_server.py | 101 +++++++++-- tests/test_sampler.py | 15 ++ 6 files changed, 343 insertions(+), 95 deletions(-) diff --git a/src/sparsevllm/engine/scheduler.py b/src/sparsevllm/engine/scheduler.py index a5d198c1..f27a4422 100644 --- a/src/sparsevllm/engine/scheduler.py +++ b/src/sparsevllm/engine/scheduler.py @@ -11,6 +11,7 @@ ) from sparsevllm.engine.sequence import Sequence, SequenceStatus from sparsevllm.engine.runtime_state import MemoryOracle +from sparsevllm.sampling_params import resolve_eos_token_ids from sparsevllm.utils.log import logger @@ -38,10 +39,12 @@ def __init__( self.chunk_prefill_size = config.chunk_prefill_size self.prefill_schedule_policy = config.prefill_schedule_policy self.eos = config.eos - configured_eos = tuple(int(token_id) for token_id in getattr(config, "eos_token_ids", ()) or ()) - if not configured_eos and int(self.eos) >= 0: - configured_eos = (int(self.eos),) - self.eos_token_ids = frozenset(configured_eos) + self.eos_token_ids = resolve_eos_token_ids( + configured_eos_token_ids=getattr( + config, "eos_token_ids", () + ), + fallback_eos_token_id=self.eos, + ) self.num_sink_tokens = config.num_sink_tokens self.num_recent_tokens = config.num_recent_tokens @@ -831,7 +834,10 @@ def postprocess( # 记录模型生成的第一个 Token seq.append_token(token_id, token_logprob, top_logprob) # 检查是否命中结束条件 - request_eos = frozenset(seq.eos_token_ids) or self.eos_token_ids + request_eos = resolve_eos_token_ids( + seq.eos_token_ids, + self.eos_token_ids, + ) if (not seq.ignore_eos and token_id in request_eos) or seq.num_completion_tokens == seq.max_tokens: seq.status = SequenceStatus.FINISHED self.decoding.remove(seq) @@ -852,7 +858,10 @@ def postprocess( ) continue seq.append_token(token_id, token_logprob, top_logprob) - request_eos = frozenset(seq.eos_token_ids) or self.eos_token_ids + request_eos = resolve_eos_token_ids( + seq.eos_token_ids, + self.eos_token_ids, + ) if (not seq.ignore_eos and token_id in request_eos) or seq.num_completion_tokens == seq.max_tokens: seq.status = SequenceStatus.FINISHED if seq in self.decoding: diff --git a/src/sparsevllm/entrypoints/openai/detokenizer.py b/src/sparsevllm/entrypoints/openai/detokenizer.py index 21546071..3de30c54 100644 --- a/src/sparsevllm/entrypoints/openai/detokenizer.py +++ b/src/sparsevllm/entrypoints/openai/detokenizer.py @@ -1,3 +1,4 @@ +from bisect import bisect_right from dataclasses import dataclass from typing import Any @@ -18,6 +19,38 @@ class DecodedFinal: raw_text_delta: str +@dataclass(frozen=True) +class _VisibleFragment: + visible_start: int + text: str + raw_start: int + raw_text: str + + @property + def visible_end(self) -> int: + return self.visible_start + len(self.text) + + def raw_offset(self, visible_offset: int) -> int: + relative = visible_offset - self.visible_start + if not 0 <= relative < len(self.text): + raise ValueError( + f"Visible offset {visible_offset} is outside fragment " + f"[{self.visible_start}, {self.visible_end})." + ) + if self.text == self.raw_text: + return self.raw_start + relative + if not self.raw_text: + return self.raw_start + + first = self.raw_text.find(self.text) + if first >= 0 and self.raw_text.find(self.text, first + 1) < 0: + return self.raw_start + first + relative + raise RuntimeError( + "Visible and raw decode fragments cannot be aligned: " + f"visible={self.text!r} raw={self.raw_text!r}." + ) + + class IncrementalDetokenizer: def __init__(self, tokenizer: Any): backend_tokenizer = getattr(tokenizer, "backend_tokenizer", None) @@ -33,8 +66,52 @@ def __init__(self, tokenizer: Any): self.token_ids: list[int] = [] self.text = "" self.raw_text = "" + self._visible_fragments: list[_VisibleFragment] = [] + self._visible_fragment_ends: list[int] = [] self.finished = False + def _record_fragment(self, text: str, raw_text: str) -> None: + if not text: + return + fragment = _VisibleFragment( + visible_start=len(self.text), + text=text, + raw_start=len(self.raw_text), + raw_text=raw_text, + ) + self._visible_fragments.append(fragment) + self._visible_fragment_ends.append(fragment.visible_end) + + def raw_offset_for_visible_prefix( + self, + visible_text_len: int, + *, + raw_text_limit: int | None = None, + ) -> int: + if not 0 <= visible_text_len <= len(self.text): + raise ValueError( + f"Visible prefix length {visible_text_len} is outside " + f"[0, {len(self.text)}]." + ) + if raw_text_limit is None: + raw_text_limit = len(self.raw_text) + if not 0 <= raw_text_limit <= len(self.raw_text): + raise ValueError( + f"Raw text limit {raw_text_limit} is outside " + f"[0, {len(self.raw_text)}]." + ) + + fragment_index = bisect_right( + self._visible_fragment_ends, + visible_text_len, + ) + if fragment_index == len(self._visible_fragments): + return raw_text_limit + raw_offset = self._visible_fragments[fragment_index].raw_offset( + visible_text_len + ) + return min(raw_offset, raw_text_limit) + def push(self, token_ids: list[int]) -> DecodedDelta: if self.finished: raise RuntimeError("Cannot push token IDs after incremental detokenization finished.") @@ -44,17 +121,24 @@ def push(self, token_ids: list[int]) -> DecodedDelta: for token_id in token_ids: token_id = int(token_id) self.token_ids.append(token_id) - text = self.visible_stream.step(self.backend_tokenizer, token_id) - raw_text = self.raw_stream.step(self.backend_tokenizer, token_id) + text = ( + self.visible_stream.step(self.backend_tokenizer, token_id) + or "" + ) + raw_text = ( + self.raw_stream.step(self.backend_tokenizer, token_id) + or "" + ) + self._record_fragment(text, raw_text) if text: text_parts.append(text) if raw_text: raw_text_parts.append(raw_text) + self.text += text + self.raw_text += raw_text text_delta = "".join(text_parts) raw_text_delta = "".join(raw_text_parts) - self.text += text_delta - self.raw_text += raw_text_delta return DecodedDelta(text=text_delta, raw_text=raw_text_delta) def finish(self, token_ids: list[int]) -> DecodedFinal: @@ -88,8 +172,11 @@ def finish(self, token_ids: list[int]) -> DecodedFinal: f"incremental={self.raw_text!r} final={final_raw_text!r}." ) - text_delta = pushed_text_delta + final_text[len(self.text):] - raw_text_delta = pushed_raw_text_delta + final_raw_text[len(self.raw_text):] + final_text_suffix = final_text[len(self.text):] + final_raw_text_suffix = final_raw_text[len(self.raw_text):] + self._record_fragment(final_text_suffix, final_raw_text_suffix) + text_delta = pushed_text_delta + final_text_suffix + raw_text_delta = pushed_raw_text_delta + final_raw_text_suffix self.text = final_text self.raw_text = final_raw_text self.finished = True diff --git a/src/sparsevllm/entrypoints/openai/dispatcher.py b/src/sparsevllm/entrypoints/openai/dispatcher.py index 351c651c..35d2b166 100644 --- a/src/sparsevllm/entrypoints/openai/dispatcher.py +++ b/src/sparsevllm/entrypoints/openai/dispatcher.py @@ -13,6 +13,7 @@ from sparsevllm.entrypoints.openai.sampling import _safe_stream_text_len from sparsevllm.llm import LLM from sparsevllm.sampling_params import SamplingParams +from sparsevllm.sampling_params import resolve_eos_token_ids from sparsevllm.utils.log import logger @@ -62,6 +63,7 @@ class _ActiveRequest: ignore_eos: bool = False terminal: threading.Event = field(default_factory=threading.Event) emitted_text_len: int = 0 + emitted_raw_text_len: int = 0 pending_token_ids: list[int] = field(default_factory=list) pending_token_logprobs: list[float | None] = field(default_factory=list) pending_top_logprobs: list[dict[int, float] | None] = field(default_factory=list) @@ -584,23 +586,14 @@ def _admit(self, item: _QueuedRequest, active: dict[int, _ActiveRequest]): if isinstance(item.prompt, list) else self.engine.tokenizer.encode(item.prompt) ) - requested_eos = getattr( - item.sampling_params, "eos_token_ids", () - ) or () - if isinstance(requested_eos, int): - requested_eos = (requested_eos,) - eos_token_ids = frozenset( - int(token_id) for token_id in requested_eos + engine_config = getattr(self.engine, "config", None) + eos_token_ids = resolve_eos_token_ids( + getattr(item.sampling_params, "eos_token_ids", ()), + getattr(engine_config, "eos_token_ids", ()), + fallback_eos_token_id=getattr( + engine_config, "eos", -1 + ), ) - if not eos_token_ids: - configured_eos = getattr( - getattr(self.engine, "config", None), - "eos_token_ids", - (), - ) or () - eos_token_ids = frozenset( - int(token_id) for token_id in configured_eos - ) active[seq_id] = _ActiveRequest( index=item.index, loop=item.loop, @@ -675,6 +668,7 @@ def _response_parser_raw_text( request: _ActiveRequest, token_ids: list[int], raw_text: str, + visible_text_len: int, ) -> str: if ( request.ignore_eos @@ -682,27 +676,32 @@ def _response_parser_raw_text( or not token_ids or int(token_ids[-1]) not in request.eos_token_ids ): - return raw_text - - content_end = len(token_ids) - while ( - content_end > 0 - and int(token_ids[content_end - 1]) in request.eos_token_ids - ): - content_end -= 1 - return request.detokenizer.tokenizer.decode( - token_ids[:content_end], - skip_special_tokens=False, - ) + parser_raw_text = raw_text + else: + content_end = len(token_ids) + while ( + content_end > 0 + and int(token_ids[content_end - 1]) + in request.eos_token_ids + ): + content_end -= 1 + parser_raw_text = request.detokenizer.tokenizer.decode( + token_ids[:content_end], + skip_special_tokens=False, + ) - @staticmethod - def _response_parser_delta(previous: str, current: str) -> str: - if not current.startswith(previous): + if not raw_text.startswith(parser_raw_text): raise RuntimeError( - "Response-parser text is not append-only: " - f"previous={previous!r} current={current!r}." + "Response-parser raw text is not a prefix of detokenized " + f"raw text: parser={parser_raw_text!r} raw={raw_text!r}." ) - return current[len(previous):] + if not request.stop: + return parser_raw_text + raw_text_len = request.detokenizer.raw_offset_for_visible_prefix( + visible_text_len, + raw_text_limit=len(parser_raw_text), + ) + return parser_raw_text[:raw_text_len] def _publish_token_deltas(self, active: dict[int, _ActiveRequest]): logprob_outputs = { @@ -721,11 +720,6 @@ def _publish_token_deltas(self, active: dict[int, _ActiveRequest]): seq_id, ([None] * len(token_ids), [None] * len(token_ids)), ) - previous_parser_raw_text = self._response_parser_raw_text( - request, - request.completion_token_ids, - request.detokenizer.raw_text, - ) request.completion_token_ids.extend(token_ids) request.completion_token_logprobs.extend(token_logprobs) request.completion_top_logprobs.extend(top_logprobs) @@ -733,15 +727,6 @@ def _publish_token_deltas(self, active: dict[int, _ActiveRequest]): request.pending_token_logprobs.extend(token_logprobs) request.pending_top_logprobs.extend(top_logprobs) request.detokenizer.push(token_ids) - parser_raw_text = self._response_parser_raw_text( - request, - request.completion_token_ids, - request.detokenizer.raw_text, - ) - raw_text_delta = self._response_parser_delta( - previous_parser_raw_text, - parser_raw_text, - ) full_text = request.detokenizer.text stop_index = _find_stop_index(full_text, request.stop) visible_text = full_text if stop_index is None else full_text[:stop_index] @@ -750,10 +735,40 @@ def _publish_token_deltas(self, active: dict[int, _ActiveRequest]): if stop_index is not None else _safe_stream_text_len(visible_text, request.stop) ) + parser_raw_text = self._response_parser_raw_text( + request, + request.completion_token_ids, + request.detokenizer.raw_text, + emit_len, + ) + if len(parser_raw_text) < request.emitted_raw_text_len: + raise RuntimeError( + "Response-parser boundary precedes emitted text: " + f"emitted={request.emitted_raw_text_len} " + f"boundary={len(parser_raw_text)}." + ) + raw_text_delta = parser_raw_text[ + request.emitted_raw_text_len: + ] text = visible_text[request.emitted_text_len:emit_len] request.emitted_text_len = emit_len if stop_index is not None: final = request.detokenizer.finish(request.completion_token_ids) + final_stop_index = _find_stop_index( + final.text, + request.stop, + ) + if final_stop_index is None: + raise RuntimeError( + "Incremental stop match disappeared during final " + "detokenization." + ) + final_parser_raw_text = self._response_parser_raw_text( + request, + request.completion_token_ids, + final.raw_text, + final_stop_index, + ) try: self.engine.abort_request( seq_id, disposition="invalidate" @@ -763,9 +778,18 @@ def _publish_token_deltas(self, active: dict[int, _ActiveRequest]): if request.chain_id is not None: request.chain_status = "invalidated" self._refresh_routing_snapshots() - if text or (stop_index is None and raw_text_delta): + if text or raw_text_delta: self._publish_pending_token_event(request, text, raw_text_delta) if stop_index is not None: + if ( + len(final_parser_raw_text) + < request.emitted_raw_text_len + ): + raise RuntimeError( + "Final response-parser text is shorter than its " + f"stream: emitted={request.emitted_raw_text_len} " + f"final={len(final_parser_raw_text)}." + ) active.pop(seq_id, None) self._mark_request_terminal(seq_id, request) self._put( @@ -774,11 +798,7 @@ def _publish_token_deltas(self, active: dict[int, _ActiveRequest]): "type": "final", "index": request.index, "text": visible_text, - "raw_text": self._response_parser_raw_text( - request, - request.completion_token_ids, - final.raw_text, - ), + "raw_text": final_parser_raw_text, "text_delta": visible_text[request.emitted_text_len:], "finish_reason": "stop", "prompt_tokens": ( @@ -822,6 +842,7 @@ def _publish_pending_token_event( request.pending_token_ids.clear() request.pending_token_logprobs.clear() request.pending_top_logprobs.clear() + request.emitted_raw_text_len += len(raw_text_delta) def _mark_request_terminal( self, @@ -1043,21 +1064,7 @@ def _publish_finished( if request is None: continue observed = len(request.completion_token_ids) - previous_parser_raw_text = self._response_parser_raw_text( - request, - request.completion_token_ids, - request.detokenizer.raw_text, - ) final = request.detokenizer.finish(completion_token_ids) - parser_raw_text = self._response_parser_raw_text( - request, - completion_token_ids, - final.raw_text, - ) - parser_raw_text_delta = self._response_parser_delta( - previous_parser_raw_text, - parser_raw_text, - ) request.pending_token_ids.extend(completion_token_ids[observed:]) request.pending_token_logprobs.extend(token_logprobs[observed:]) request.pending_top_logprobs.extend(top_logprobs[observed:]) @@ -1065,9 +1072,35 @@ def _publish_finished( request.completion_token_ids = list(completion_token_ids) request.completion_token_logprobs = list(token_logprobs) request.completion_top_logprobs = list(top_logprobs) - finish_reason = "length" if len(completion_token_ids) >= request.max_tokens else "stop" text = final.text stop_index = _find_stop_index(text, request.stop) + parser_raw_text = self._response_parser_raw_text( + request, + completion_token_ids, + final.raw_text, + stop_index if stop_index is not None else len(text), + ) + if len(parser_raw_text) < request.emitted_raw_text_len: + raise RuntimeError( + "Final response-parser text is shorter than its stream: " + f"emitted={request.emitted_raw_text_len} " + f"final={len(parser_raw_text)}." + ) + parser_raw_text_delta = parser_raw_text[ + request.emitted_raw_text_len: + ] + ended_by_eos = ( + not request.ignore_eos + and bool(completion_token_ids) + and int(completion_token_ids[-1]) + in request.eos_token_ids + ) + finish_reason = ( + "stop" + if ended_by_eos + or len(completion_token_ids) < request.max_tokens + else "length" + ) if stop_index is not None: text = text[:stop_index] finish_reason = "stop" diff --git a/src/sparsevllm/sampling_params.py b/src/sparsevllm/sampling_params.py index f06b3ea7..dc8d425c 100644 --- a/src/sparsevllm/sampling_params.py +++ b/src/sparsevllm/sampling_params.py @@ -1,6 +1,37 @@ +from collections.abc import Iterable from dataclasses import dataclass +def _as_eos_token_id_set( + value: int | Iterable[int] | None, +) -> frozenset[int]: + if value is None: + return frozenset() + if isinstance(value, int): + return frozenset({int(value)}) + return frozenset(int(token_id) for token_id in value) + + +def resolve_eos_token_ids( + request_eos_token_ids: int | Iterable[int] | None = (), + configured_eos_token_ids: int | Iterable[int] | None = (), + *, + fallback_eos_token_id: int | None = -1, +) -> frozenset[int]: + requested = _as_eos_token_id_set(request_eos_token_ids) + if requested: + return requested + configured = _as_eos_token_id_set(configured_eos_token_ids) + if configured: + return configured + if ( + fallback_eos_token_id is not None + and int(fallback_eos_token_id) >= 0 + ): + return frozenset({int(fallback_eos_token_id)}) + return frozenset() + + @dataclass class SamplingParams: temperature: float = 1.0 diff --git a/tests/test_openai_api_server.py b/tests/test_openai_api_server.py index e9b6cca8..72b65563 100644 --- a/tests/test_openai_api_server.py +++ b/tests/test_openai_api_server.py @@ -83,13 +83,16 @@ def _byte_level_tokenizer(*, special_tokens=()): return _FastTokenizerAdapter(tokenizer) -async def _dispatcher_items_for_text(text, *, stop=()): +async def _dispatcher_items_for_token_ids( + tokenizer, + token_ids, + *, + stop=(), + incremental=True, +): from sparsevllm.entrypoints.openai.api_server import AsyncEngineDispatcher, _ActiveRequest from sparsevllm.entrypoints.openai.detokenizer import IncrementalDetokenizer - tokenizer = _byte_level_tokenizer() - token_ids = tokenizer.encode(text) - class Engine: def __init__(self): self.tokenizer = tokenizer @@ -122,15 +125,16 @@ def exit(self): } items = [] try: - for token_id in token_ids: - engine.last_step_token_outputs = [(7, [token_id])] - engine.last_step_logprob_outputs = [(7, [None], [None])] - dispatcher._publish_token_deltas(active) - await asyncio.sleep(0) - while not output_queue.empty(): - items.append(output_queue.get_nowait()) - if 7 not in active: - break + if incremental: + for token_id in token_ids: + engine.last_step_token_outputs = [(7, [token_id])] + engine.last_step_logprob_outputs = [(7, [None], [None])] + dispatcher._publish_token_deltas(active) + await asyncio.sleep(0) + while not output_queue.empty(): + items.append(output_queue.get_nowait()) + if 7 not in active: + break if 7 in active: dispatcher._publish_finished( @@ -145,6 +149,15 @@ def exit(self): return items +async def _dispatcher_items_for_text(text, *, stop=()): + tokenizer = _byte_level_tokenizer() + return await _dispatcher_items_for_token_ids( + tokenizer, + tokenizer.encode(text), + stop=stop, + ) + + class _TestRequest: def __init__(self, app): self.app = app @@ -2606,6 +2619,7 @@ def exit(self): completion_top_logprobs=[None], detokenizer=detokenizer, emitted_text_len=1, + emitted_raw_text_len=1, ) } try: @@ -2645,7 +2659,7 @@ def exit(self): loop=asyncio.get_running_loop(), output_queue=output_queue, prompt_token_ids=[10], - max_tokens=len(completion_token_ids) + 1, + max_tokens=len(completion_token_ids), stop=[], completion_token_ids=[], completion_token_logprobs=[], @@ -2688,6 +2702,7 @@ def exit(self): self.assertEqual(token_item["raw_text_delta"], "Paris") self.assertEqual(final_item["raw_text"], "Paris") + self.assertEqual(final_item["finish_reason"], "stop") self.assertEqual(final_item["token_ids"], completion_token_ids) self.assertEqual( final_item["completion_tokens"], @@ -2759,6 +2774,7 @@ def exit(self): self.assertEqual(token_item["raw_text_delta"], "Paris") self.assertEqual(final_item["raw_text"], "Paris") + self.assertEqual(final_item["finish_reason"], "length") async def test_dispatcher_streams_complete_unicode_with_pending_logprobs(self): from sparsevllm.entrypoints.openai.api_server import AsyncEngineDispatcher, _ActiveRequest @@ -3265,6 +3281,7 @@ def exit(self): dispatcher._publish_token_deltas(active) token_item = await asyncio.wait_for(output_queue.get(), timeout=1) self.assertEqual(token_item["text"], "a") + self.assertEqual(token_item["raw_text_delta"], "a") engine.last_step_token_outputs = [(7, token_ids[2:])] engine.last_step_logprob_outputs = [ (7, [None] * len(token_ids[2:]), [None] * len(token_ids[2:])) @@ -3276,9 +3293,57 @@ def exit(self): self.assertEqual(final_item["type"], "final") self.assertEqual(final_item["text"], "a") + self.assertEqual(final_item["raw_text"], "a") self.assertEqual(final_item["text_delta"], "") self.assertEqual(engine.aborted, [7]) + async def test_dispatcher_maps_visible_stop_across_special_tokens(self): + tokenizer = _byte_level_tokenizer( + special_tokens=[""], + ) + special_token_id = tokenizer._tokenizer.token_to_id("") + cases = [ + ( + [special_token_id] + tokenizer.encode("special"), + "special", + "", + "", + ), + ( + tokenizer.encode("okST") + + [special_token_id] + + tokenizer.encode("OP"), + "STOP", + "ok", + "ok", + ), + ] + + for incremental in (True, False): + for token_ids, stop, expected_text, expected_raw in cases: + with self.subTest( + incremental=incremental, + stop=stop, + ): + items = await _dispatcher_items_for_token_ids( + tokenizer, + token_ids, + stop=[stop], + incremental=incremental, + ) + streamed_raw = "".join( + item["raw_text_delta"] + for item in items + if item["type"] == "token" + ) + final = next( + item for item in items if item["type"] == "final" + ) + + self.assertEqual(streamed_raw, expected_raw) + self.assertEqual(final["text"], expected_text) + self.assertEqual(final["raw_text"], expected_raw) + async def test_chat_completion_response_shape(self): from sparsevllm.entrypoints.openai.api_server import RequestHandle, _chat_completion_response @@ -3686,6 +3751,7 @@ def exit(self): final = next(item for item in items if item["type"] == "final") self.assertEqual(final["text"], "answer") + self.assertEqual(final["raw_text"], "answer") self.assertEqual(final["chain_status"], "invalidated") self.assertEqual(engine.aborted, [(7, "invalidate")]) @@ -3781,6 +3847,12 @@ class Admission: prefilled_tokens = 1 class Engine: + config = type( + "Config", + (), + {"eos_token_ids": (41, 42), "eos": -1}, + )() + def __init__(self): self.tokenizer = tokenizer @@ -3815,6 +3887,7 @@ def admit_request(self, *_args, **_kwargs): dispatcher._admit(item, active) await asyncio.wait_for(future, timeout=1) + self.assertEqual(active[7].eos_token_ids, frozenset({41, 42})) dispatcher._publish_finished(active, [(7, [], [], [])]) final = await asyncio.wait_for(output_queue.get(), timeout=1) diff --git a/tests/test_sampler.py b/tests/test_sampler.py index 87126c37..62e599a8 100644 --- a/tests/test_sampler.py +++ b/tests/test_sampler.py @@ -8,6 +8,7 @@ from sparsevllm.engine.sequence import Sequence from sparsevllm.layers.sampler import Sampler from sparsevllm.sampling_params import SamplingParams +from sparsevllm.sampling_params import resolve_eos_token_ids class SamplerTest(unittest.TestCase): @@ -85,6 +86,20 @@ def test_sampling_penalty_defaults_are_neutral(self): ) self.assertIs(output, logits) + def test_resolve_eos_token_ids_uses_request_precedence(self): + self.assertEqual( + resolve_eos_token_ids((7, 8), (9,), fallback_eos_token_id=10), + frozenset({7, 8}), + ) + self.assertEqual( + resolve_eos_token_ids((), (9,), fallback_eos_token_id=10), + frozenset({9}), + ) + self.assertEqual( + resolve_eos_token_ids((), (), fallback_eos_token_id=10), + frozenset({10}), + ) + def test_presence_and_repetition_penalty_formulas(self): sampler = Sampler() logits = torch.tensor( From 5436b37f05ce999c1d33217fac02c9276745dea7 Mon Sep 17 00:00:00 2001 From: Codex Date: Wed, 5 Aug 2026 01:28:35 +0800 Subject: [PATCH 14/48] docs: document glm4 support --- README.md | 1 + docs/en/features/supported-models.md | 7 + docs/zh/design/README.md | 1 + docs/zh/design/glm-4.7-flash-support-plan.md | 1047 ++++++++++++++++++ docs/zh/features/supported-models.md | 7 + 5 files changed, 1063 insertions(+) create mode 100644 docs/zh/design/glm-4.7-flash-support-plan.md diff --git a/README.md b/README.md index 39f8b2f0..74db0d7f 100644 --- a/README.md +++ b/README.md @@ -70,6 +70,7 @@ Read the method overview and integration rules in | Qwen3MoE | ✅ | | Qwen3.5 / Qwen3.6 | ✅ | | Qwen3.5 / Qwen3.6 MoE | ✅ | +| GLM-4.7-Flash | ✅ | | Llama 3 / 3.1 | ✅ | | MiniMax M2.7 | ✅ | diff --git a/docs/en/features/supported-models.md b/docs/en/features/supported-models.md index 8ff710a8..c2b50f7f 100644 --- a/docs/en/features/supported-models.md +++ b/docs/en/features/supported-models.md @@ -17,6 +17,7 @@ parallel size must use that value. | Qwen3MoE | `qwen3_moe` | BF16 / FP16 / block FP8 | ✅ (TP > 1: BF16 model dtype only) | 1 only | ✅ | | Qwen3.5 / Qwen3.6 | `qwen3_5` | BF16 / block FP8 | ✅ | 1 only | 1 only | | Qwen3.6 MoE | `qwen3_5_moe` | BF16 / block FP8 | ✅ | 1 only | ✅ | +| GLM-4.7-Flash | `glm4_moe_lite` | BF16 | 1 only | 1 only | 1 only | | Llama 3 / 3.1 | `llama` | BF16 / FP16 | ✅ | 1 only | 1 only | | MiniMax M2.7 | `minimax_m2` | block FP8 with BF16 non-quantized weights | ✅ | 1 only | ✅ | @@ -37,6 +38,11 @@ a `128 x 128` weight block size. Qwen3.5/Qwen3.6 dense configurations are normalized internally to `model_type=qwen3_5`; Qwen3.6 MoE uses `model_type=qwen3_5_moe`. +GLM-4.7-Flash support is currently limited to BF16, eager execution, vanilla +latent MLA, and `TP=DP=EP=1`. The loader intentionally skips the checkpoint's +MTP layer. TP greater than 1, expert parallelism, CUDA Graph, prefix caching, +quantization, and sparse methods remain unsupported until separately validated. + ## Sparse Method Support | Model | Vanilla | StreamingLLM | SnapKV | H2O | PyramidKV | OmniKV | QuEST | R-KV | SkipKV | DeltaKV | @@ -46,6 +52,7 @@ normalized internally to `model_type=qwen3_5`; Qwen3.6 MoE uses | Qwen3MoE | ✅ | ✅ | ✅ | Experimental⁴ | ✅ | ✅ | ✅ | ✅ | — | — | | Qwen3.5 / Qwen3.6 | ✅ | ✅ | ✅ | Experimental⁴ | ✅ | ✅ | ✅ | ✅ | — | Matched checkpoint³ | | Qwen3.6 MoE | ✅ | ✅ | ✅ | Experimental⁴ | ✅ | ✅ | ✅ | ✅ | — | — | +| GLM-4.7-Flash | ✅ | — | — | — | — | — | — | — | — | — | | Llama 3 / 3.1 | ✅ | ✅ | ✅ | Experimental⁴ | ✅ | ✅ | ✅ | ✅ | Selected checkpoint¹ | Compressor required² | | MiniMax M2.7 | ✅ | ✅ | ✅ | Experimental⁴ | ✅ | ✅ | ✅ | ✅ | — | — | diff --git a/docs/zh/design/README.md b/docs/zh/design/README.md index ac2c815a..04a9eb1f 100644 --- a/docs/zh/design/README.md +++ b/docs/zh/design/README.md @@ -4,3 +4,4 @@ - [架构](architecture.md) - [Sparse-vLLM 控制图](control-map.md) +- [GLM-4.7-Flash 支持实施计划](glm-4.7-flash-support-plan.md) diff --git a/docs/zh/design/glm-4.7-flash-support-plan.md b/docs/zh/design/glm-4.7-flash-support-plan.md new file mode 100644 index 00000000..f1b50b62 --- /dev/null +++ b/docs/zh/design/glm-4.7-flash-support-plan.md @@ -0,0 +1,1047 @@ +# GLM-4.7-Flash 支持实施计划 + +- 状态：Implemented（TP1 已验证；TP4/EP1 待验证） +- 更新时间：2026-08-04 +- 目标分支：codex/glm-4.7-flash +- 目标模型：zai-org/GLM-4.7-Flash / glm4_moe_lite + +## 1. 摘要 + +本计划为 Sparse-vLLM 增加 GLM-4.7-Flash 的第一阶段支持。实现以 latent MLA 为基础，不通过持久化展开式 K/V 获得临时兼容，也不把 GLM 特例塞进现有显式 K/V 接口。 + +核心决策如下： + +1. 从固定版本的 LightLLM vendor 最小 MLA Triton 数学核心，移除 LightLLM runtime 依赖，并先以独立 Torch oracle 验证。 +2. 将当前 PrefillComputeView / DecodeComputeView 从固定双 K/V 改为“公共 metadata + tagged payload”。 +3. 保留 CacheManager 对 slot、请求生命周期和 logical view 的所有权；由 storage strategy 管理显式 K/V 或 MLA latent 的物理张量。 +4. 模型只表达 GLM 投影、部分 RoPE、K/V absorption、MoE/router/shared-expert 等语义；模型不选择或直接导入具体 kernel。 +5. MLA kernel 通过 OpSpec -> OpResolver -> Provider -> kernel 路径绑定，初始化后不做运行时静默 fallback。 +6. MoE 只抽取 Qwen3-MoE 中真正通用的 packed-expert 物理执行部分，不构建万能 MoE 模型框架。 +7. GLM response parsing 继续委托 Transformers；本地只提供缺失的声明式 response template。 + +首版目标是 BF16、eager、vanilla latent MLA、TP1 和 TP4/EP1 的短上下文正确性。CUDA Graph、Prefix Cache、稀疏方法、MTP、量化和实用长上下文能力均需后续单独验证。 + +## 2. 目标与非目标 + +### 2.1 目标 + +- 正确加载 Glm4MoeLiteForCausalLM 的基础模型权重。 +- 正确执行层 0–46；明确识别并跳过 checkpoint 中单独的 MTP 层 47。 +- 使用 latent MLA cache：每层每 token 持久化 512 维 latent 和 64 维 RoPE key。 +- 支持 chunked prefill 的完整历史可见性。 +- 支持 padded static/eager decode，不发生 slot=-1 越界写。 +- 支持 GLM 的首层 Dense、后续 routed MoE、shared expert 和 biased-sigmoid routing。 +- 复用当前 TP/EP、packed expert、provider registry、linear、RMSNorm 和 SwiGLU 基础设施。 +- 提供可复现、artifact-backed 的 tiny 与真实模型验证。 + +### 2.2 非目标 + +首版不包含： + +- MTP/speculative decoding。 +- 非 BF16 checkpoint 或量化权重。 +- CUDA Graph。 +- Prefix Cache 或 prefix offload。 +- SnapKV、PyramidKV、OmniKV、QuEST、H2O、RKV、StreamingLLM、DeltaKV 等稀疏方法。 +- EP 大于 1 的真实模型支持声明。 +- 128K/202K 长上下文吞吐或显存能力声明。 +- LightLLM Python package 或 runtime 依赖。 +- 在 kernel 失败后切换到 Torch/其他 backend 的运行时 fallback。 + +## 3. 已知模型契约 + +官方配置来源： + +- https://huggingface.co/zai-org/GLM-4.7-Flash/blob/main/config.json +- https://github.com/huggingface/transformers/blob/main/src/transformers/models/glm4_moe_lite/modeling_glm4_moe_lite.py + +关键结构： + +| 项目 | 值 | +| --- | ---: | +| model_type | glm4_moe_lite | +| 基础层数 | 47 | +| MTP 层数 | 1 | +| hidden size | 2048 | +| attention heads | 20 | +| KV heads | 20（展开语义） | +| q LoRA rank | 768 | +| kv LoRA rank | 512 | +| qk no-PE dim | 192 | +| qk RoPE dim | 64 | +| qk head dim | 256 | +| value head dim | 256 | +| routed experts | 64 | +| experts per token | 4 | +| shared experts | 1 | +| MoE intermediate size | 1536 | +| Dense intermediate size | 10240 | +| routed scaling factor | 1.8 | +| first Dense layers | 1 | +| dtype | BF16 | + +必须区分以下维度： + +~~~ +kv_lora_rank = 512 +rope_dim = 64 +latent_cache_width = 576 +qk_head_dim = 192 + 64 = 256 +softmax_scale = 256 ** -0.5 +~~~ + +不能从 latent cache width 576 推导 attention scale，也不能使用 Transformers 映射后的 config.head_dim=64 作为 Sparse-vLLM cache head dimension。 + +## 4. 当前运行时约束 + +当前实现从抽象层开始假定显式双 K/V： + +- src/sparsevllm/engine/cache_manager/base.py 的 PrefillComputeView 和 DecodeComputeView 固定包含 k_cache / v_cache。 +- CacheManager 使用 num_key_value_heads 和 head_dim 推导物理 cache shape 与容量。 +- StandardCacheManager 直接分配 K/V tensor，并通过 slot table 暴露计算视图。 +- TritonAttentionBackend 直接消费 view.k_cache / view.v_cache。 +- eager decode 仍可能经过 static batch bucket padding；padding sequence 的 slot_mapping 可以是 -1。 +- chunked prefill 的 view 表示完整可见上下文，而不仅是当前输入 chunk。 + +因此禁止以下临时方案： + +- 把 latent 和 RoPE tensor 假装成 k_cache / v_cache。 +- 在 ComputeView 中增加一组互斥的 optional tensor 字段。 +- 通过未校验的 metadata 字典把 MLA tensor 旁路给 kernel。 +- 只展开当前 prefill chunk，忽略已缓存历史。 +- 原样复制 LightLLM 的 infer_state wrapper。 + +## 5. 目标架构 + +~~~ +Glm4MoeLiteForCausalLM + | + +-- Glm4MoeLiteAttention + | | + | +-- q_a/q_b, kv_a/kv_b, partial RoPE, K/V absorption + | | + | +-- MLAAttention (semantic layer) + | | + | +-- AttentionComputeView[AttentionPayload] + | | +-- AttentionViewMeta + | | +-- MlaLatentPayload + | | + | +-- MlaAttentionOpSpec + | +-- OpResolver(DeviceCaps) + | +-- MlaTritonProvider + | +-- vendored LightLLM-derived Triton kernels + | + +-- Glm4MoeLiteMlp + +-- Dense Qwen3MLP (layer 0) + +-- GlmBiasedSigmoidRouter + +-- shared PackedMoeExperts + +-- Qwen3MLP shared expert + +CacheManager + +-- slot allocator / request lifetime / logical selection + +-- ExplicitKVStorage | MlaLatentStorage +~~~ + +所有权原则： + +- 模型拥有模型语义和逻辑权重。 +- CacheManager 拥有 slot、请求生命周期和 logical visibility。 +- Storage 拥有物理 cache tensor、store 和内存核算。 +- ComputeView 描述一次 attention 计算可见的 metadata 与 payload。 +- Provider 拥有 kernel-specific layout、workspace、launch config 和调用。 +- Platform / DeviceCaps 拥有设备发现和能力事实。 + +## 6. ComputeView 改造 + +### 6.1 设计 + +建议引入： + +~~~python +@dataclass(frozen=True) +class AttentionViewMeta: + active_slots: torch.Tensor + req_indices: torch.Tensor + context_lens: torch.Tensor + max_context_len: int | None = None + attn_score: torch.Tensor | None = None + temp_slots: torch.Tensor | None = None + + +@dataclass(frozen=True) +class ExplicitKVPayload: + k_cache: torch.Tensor + v_cache: torch.Tensor + backend: str = "dense" + metadata: dict[str, Any] | None = None + + +@dataclass(frozen=True) +class MlaLatentPayload: + latent_cache: torch.Tensor + rope_cache: torch.Tensor + + +AttentionPayload = ExplicitKVPayload | MlaLatentPayload + + +@dataclass +class PrefillComputeView: + meta: AttentionViewMeta + payload: AttentionPayload + + +@dataclass +class DecodeComputeView: + meta: AttentionViewMeta + payload: AttentionPayload +~~~ + +Prefill 和 decode 保留不同类型，因为二者的执行、workspace 和生命周期不同；二者共用 payload 类型和公共 metadata。 + +### 6.2 不变量 + +- 不允许一个 view 同时包含显式 K/V 与 MLA latent。 +- 通用 Attention 必须显式要求 ExplicitKVPayload。 +- MLAAttention 必须显式要求 MlaLatentPayload。 +- payload 类型不匹配时在进入 kernel 前失败。 +- backend 和 backend-specific metadata 只属于 ExplicitKVPayload 或对应 provider，不属于公共 metadata。 +- AttentionViewMeta 不携带模型权重或 kernel workspace。 + +### 6.3 迁移策略 + +先进行不改变行为的机械迁移： + +~~~ +view.k_cache -> view.payload.k_cache +view.v_cache -> view.payload.v_cache +view.active_slots -> view.meta.active_slots +view.context_lens -> view.meta.context_lens +view.backend -> view.payload.backend +~~~ + +所有现有 manager 在该提交中只生成 ExplicitKVPayload。迁移提交不得同时引入 MLA 逻辑，以便将现有测试失败明确归因于接口改造。 + +## 7. Cache storage 契约 + +### 7.1 目标 + +避免创建以下类层次： + +~~~ +GlmStandardCacheManager +GlmSnapKVCacheManager +GlmQuestCacheManager +... +~~~ + +CacheManager 继续按 sparse method 选择；物理 payload 通过 composition 委托给 storage。 + +### 7.2 建议接口 + +~~~python +class AttentionCacheStorage(Protocol): + layout: CacheLayout + + def allocate( + self, + *, + num_layers: int, + num_slots: int, + device: torch.device, + ) -> None: ... + + def layer_payload(self, layer_idx: int) -> AttentionPayload: ... + + def store( + self, + layer_idx: int, + slot_mapping: torch.Tensor, + payload: Any, + ) -> None: ... + + def bytes_per_slot_per_layer(self) -> int: ... + + def accounting_tensors(self) -> tuple[torch.Tensor, ...]: ... +~~~ + +具体类： + +~~~ +ExplicitKVStorage + k_cache [layers, slots, kv_heads, head_dim] + v_cache [layers, slots, kv_heads, head_dim] + +MlaLatentStorage + latent_cache [layers, slots, 1, 512] + rope_cache [layers, slots, 1, 64] +~~~ + +### 7.3 内存核算 + +所有以下逻辑必须从 storage 获取，不再使用全局固定公式 2 * num_kv_heads * head_dim： + +- 每 slot 每 layer 字节数。 +- 最大可分配 slot 数。 +- prefix block 字节数。 +- allocator/debug tensor 汇总。 +- live-token 和 reserved-token 内存报告。 + +accounting_tensors() 必须显式返回 storage 持有的物理 tensor，不能依赖递归扫描任意对象。 + +### 7.4 Padding 安全 + +store() 和 vendored copy kernel 必须跳过 slot_mapping < 0。以下两种模式都必须覆盖： + +- 真实 request slot。 +- static/eager padded batch 的 slot=-1。 + +禁止依赖“关闭 CUDA Graph 后不会 padding”的假设。 + +## 8. LightLLM MLA kernel vendoring + +### 8.1 固定上游 + +- 仓库：https://github.com/ModelTC/lightllm +- commit：65c174ee95ac6a6fd36b18b63d0b33d97e76b770 +- license：Apache-2.0 + +本地目录： + +~~~ +src/sparsevllm/triton_kernel/mla/ +~~~ + +### 8.2 最小来源映射 + +| 上游文件 | 本地计划文件 | 处理方式 | +| --- | --- | --- | +| common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding_stage1.py | decode_stage1.py | 保留 Triton 数学核心，改为显式 tensor/stride/config API | +| common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding_stage2.py | decode_stage2.py | 保留归并数学核心，移除 LightLLM 状态依赖 | +| common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding.py | decode_schedule.py | 不原样复制；仅参考 stage 调度，本地重写 workspace/launch API | +| common/basemodel/triton_kernel/kv_copy/mla_copy_kv.py | copy_latent.py | 增加 slot=-1 mask 与本地 layout 校验 | +| models/deepseek2/triton_kernel/sample_kv.py | gather_latent.py | 适配 slot-indirected ragged/full-history gather | + +每个衍生文件必须： + +- 保留或补充 Apache-2.0 SPDX 标识。 +- 注明原始仓库、固定 commit 和原文件路径。 +- 注明本地修改。 +- 不 import lightllm。 + +目录同时包含： + +~~~ +README.md +LICENSE.lightllm +~~~ + +### 8.3 明确不复制 + +- LightLLM infer_state。 +- request manager / memory manager。 +- 动态挂载到 infer_state 的 workspace。 +- frozendict。 +- LightLLM global kernel config 或 autotuner。 +- LightLLM device utilities。 +- forward 内的动态 vSM 探测和 allocation。 +- context_flashattention_nopad_with_v.py 作为 GLM 首版 prefill kernel。 + +最后一项不能直接用于 GLM：其契约要求 q_nope_dim == v_dim，而 GLM 是 192 与 256。 + +### 8.4 Vendor kernel API + +Decode 输入： + +~~~ +q_nope_absorbed [batch, local_heads, 512] +q_rope [batch, local_heads, 64] +latent_cache [slots, 1, 512] +rope_cache [slots, 1, 64] +active_slots [rows, max_context_len] +req_indices [batch] +context_lens [batch] +softmax_scale scalar, fixed from qk_head_dim=256 +workspace explicit tensors +launch_config immutable value object +~~~ + +输出： + +~~~ +latent_output [batch, local_heads, 512] +~~~ + +Kernel 不读取 HF config、model type、环境变量或 backend 名称。 + +## 9. MLA operator/provider + +新增： + +~~~ +src/sparsevllm/operators/mla_attention.py +~~~ + +建议 spec： + +~~~python +@dataclass(frozen=True) +class MlaAttentionOpSpec: + num_q_heads: int + kv_lora_rank: int + rope_dim: int + qk_head_dim: int + value_head_dim: int + activation_dtype: torch.dtype + cache_dtype: torch.dtype + tp_size: int + cuda_graph: bool +~~~ + +softmax_scale 可以由经过校验的 qk_head_dim 得出，也可作为只读规范化字段保存；禁止从 kv_lora_rank + rope_dim 得出。 + +Provider 选择链： + +~~~ +MLAAttention semantic call + -> MlaAttentionOpSpec + -> OpResolver(DeviceCaps) + -> MlaTritonProvider + -> prepare workspace/layout + -> decode_stage1/decode_stage2 +~~~ + +约束： + +- provider 在 model construction/ModelRunner 初始化阶段绑定一次。 +- forward hot path 不 resolve provider。 +- supports(spec, caps) 检查 platform、Triton、dtype、维度、TP、graph 和已验证硬件。 +- 未验证的 GPU 架构不宣称支持。 +- 首版只有一个 production provider；Torch 只用于测试 oracle。 +- kernel/JIT 失败直接暴露，不 catch 后切换实现。 +- Provider 使用 Platform / DeviceCaps，通用代码不直接探测 torch.cuda。 + +LightLLM 固定 commit 没有目标 H100/GLM shape 的现成 tuning config，因此首版使用明确、安全的静态 launch config；只有 H100 数值和性能实测后才将该 shape 标记为已验证。 + +## 10. Prefill 数据路径 + +MLA persistent cache 只保存 latent/rope，但 chunked prefill 必须看到完整历史。 + +正确顺序： + +1. 当前 chunk 计算 q_a/q_b 和 kv_a。 +2. 对当前 chunk 的 512 维 latent 做 RMSNorm。 +3. 对当前 chunk 的 64 维 key 应用 RoPE。 +4. 通过 padding-safe store 写入持久 latent/rope cache。 +5. 根据 active_slots + req_indices + context_lens gather 每条序列的完整历史 latent/rope，包括之前 chunk 和当前 chunk。 +6. 使用 kv_b_proj 将完整历史 latent 临时展开为： + - k_nope [visible_tokens, local_heads, 192] + - v [visible_tokens, local_heads, 256] +7. 将历史 RoPE key 扩展/映射到 local heads，与 k_nope 形成 256 维 K。 +8. 使用 Sparse-vLLM 现有 contiguous/slot-compatible 256 维 prefill attention kernel。 +9. prefill 完成后释放 gather/expanded workspace，不将其写回持久 cache。 + +需要单独的临时 workset： + +~~~python +@dataclass +class MlaPrefillWorkset: + gathered_latent: torch.Tensor + gathered_rope: torch.Tensor + expanded_k: torch.Tensor + expanded_v: torch.Tensor + packed_offsets: torch.Tensor +~~~ + +Workspace 必须有显式显存预算和有界分配。无法满足预算时应在运行 kernel 前报错，不允许降级为缺失历史或缩短上下文。 + +## 11. Decode 数据路径 + +1. 当前 token 计算： + - q_nope [batch, local_heads, 192] + - q_rope [batch, local_heads, 64] + - latent [batch, 1, 512] + - k_rope [batch, 1, 64] +2. padding-safe store 将 latent/k_rope 写入当前 slot。 +3. 使用 K_b 将 query no-PE 部分吸收到 latent 空间： + + ~~~ + q_nope_absorbed [batch, local_heads, 512] + ~~~ + +4. MLA provider 直接对 MlaLatentPayload 计算 attention，输出： + + ~~~ + latent_output [batch, local_heads, 512] + ~~~ + +5. 使用 V_b 将输出恢复为： + + ~~~ + value_output [batch, local_heads, 256] + ~~~ + +6. flatten local heads，进入 row-parallel o_proj。 + +Decode 计算中的 logits 是： + +~~~ +q_nope_absorbed @ latent_cache.T ++ q_rope @ rope_cache.T +~~~ + +整体使用 256 ** -0.5 scale。 + +## 12. 模型实现 + +新增： + +~~~ +src/sparsevllm/models/glm4_moe_lite.py +~~~ + +模型文件负责： + +- token embedding / final norm / LM head。 +- q_a_proj、q_a_layernorm、q_b_proj。 +- kv_a_proj_with_mqa、kv_a_layernorm、kv_b_proj。 +- 192+64 partial/interleaved RoPE。 +- K/V absorption 和 V reconstruction。 +- Dense/MoE layer topology。 +- checkpoint 逻辑名称映射。 + +模型文件不负责： + +- import vendored kernel。 +- 选择 provider。 +- 探测 GPU。 +- 分配持久 cache。 +- 维护 slot table。 +- 运行时 backend fallback。 + +TP 布局建议： + +- q_a_proj、kv_a_proj_with_mqa 和对应 low-rank norm：replicated。 +- q_b_proj、kv_b_proj：按 local attention heads 做 column/local-head shard。 +- o_proj：row parallel。 +- 每 rank 持有本 rank 的 K_b/V_b head slices。 +- persistent latent/rope cache 在各 TP rank 上语义一致。 + +首版验证 TP1 和 TP4；TP2 可以保留配置能力，但未经实测不列为已验证。 + +## 13. MoE 复用边界 + +### 13.1 抽取内容 + +从 src/sparsevllm/models/qwen3_moe.py 抽取 packed-expert 物理执行： + +- W13/W2 storage。 +- unquantized TP intermediate shard。 +- EP local expert ownership。 +- MoeOpSpec 和 provider resolution。 +- projection loading。 +- local expert completeness validation。 + +建议形成共享 PackedMoeExperts，由 Qwen3-MoE 与 GLM 分别组合。 + +### 13.2 不抽取内容 + +首轮不抽取： + +- Qwen softmax router。 +- 整个 Qwen MoE block。 +- model-specific chunking/all-reduce/debug。 +- checkpoint regex。 +- remote-expert skip bookkeeping。 +- MiniMax packed-expert 实现。 + +这些内容保留在模型层，以避免为了两个不同模型构建过度通用的 MoE framework。 + +### 13.3 GLM MoE + +~~~ +layer 0: + Qwen3MLP(hidden=2048, intermediate=10240) + +layer 1-46: + routed = PackedMoeExperts(x, topk_ids, topk_weights) + shared = Qwen3MLP(hidden=2048, intermediate=1536)(x) + output = world_reduce(routed) + shared +~~~ + +Router 必须复现： + +1. FP32 gate linear。 +2. sigmoid。 +3. correction bias 只影响 expert selection。 +4. 从原始 sigmoid score gather selected weights。 +5. Top-4 normalization。 +6. 乘 routed_scaling_factor=1.8。 + +当前 MiniMax router kernel 硬编码 256 experts/Top-8。将其泛化为支持 GLM 64/Top-4 应作为独立提交；不复制一个 GLM 专属基础 kernel，也不在 packed-expert 抽取提交中顺带迁移 MiniMax storage。 + +## 14. 权重加载与 MTP 边界 + +Loader 要求： + +- 所有基础层 0–46 权重必须被消费。 +- 所有本地 EP experts 的 gate/up/down 权重必须完整。 +- 只允许跳过其他 EP rank 持有的 expert 权重。 +- e_score_correction_bias 保持 FP32。 +- 只允许精确跳过 model.layers.47.* MTP 权重。 +- 任何其他 missing/unexpected tensor 直接失败。 + +加载结束时记录： + +- loaded tensor count。 +- local expert shard count。 +- remote EP skipped count。 +- MTP skipped count。 +- 其他 skipped/missing/unexpected count。 + +MTP 不得通过宽泛 silent ignore 机制掩盖错误层数或命名变化。 + +## 15. Response parsing + +Sparse-vLLM 继续使用 Transformers parser： + +~~~ +tokenizer.parse_response() # 非流式 +tokenizer.get_response_parser() # 流式 +~~~ + +官方 GLM tokenizer 当前缺少 response_template，因此本地仅增加声明式 fallback，描述： + +- ...。 +- function-name ... 。 +- ...。 +- ...。 + +禁止实现另一套 XML parser、streaming state machine 或 tool-call parser。 + +CLI 建议新增通用 auto；现有 qwen3、minimax_m2 保留兼容，glm47 仅作为 alias，不选择另一套 parser 实现。 + +## 16. 首版兼容矩阵 + +| 能力 | 首版状态 | 门禁 | +| --- | --- | --- | +| BF16 | 已实现并验证 | tiny + 真实权重验证 | +| eager | 已实现并验证 | 必须显式启用 | +| vanilla latent MLA | 已实现并验证 | kernel/storage/model 全链路 | +| TP1 | 已实现并验证 | 真实模型 | +| TP4/EP1 | 已实现路径，尚未验证 | 验证前不写入支持声明 | +| EP > 1 | 暂不支持 | 后续独立验证 | +| CUDA Graph | 暂不支持 | 配置阶段拒绝 | +| Prefix Cache/offload | 暂不支持 | 配置阶段拒绝 | +| sparse methods | 暂不支持 | method registry 显式拒绝 | +| MTP | 暂不支持 | loader 精确跳过 | +| quantization | 暂不支持 | checkpoint config 拒绝 | +| 128K/202K | 未验证 | 不写入支持声明 | + +配置校验必须发生在模型构造和 cache 分配之前，给出具体不兼容原因。 + +## 17. 分阶段执行计划 + +### 阶段 0：同步与基线 + +工作： + +- 同步当前分支与目标 main 基线。 +- 确认 worktree 干净。 +- 记录 Python、Torch、Transformers、Triton、CUDA 和 GPU 信息。 +- 校验本地 checkpoint 是否具备 48 个 shard、index 和可追踪 revision/hash。 +- 将 Hugging Face、Xet、Triton、TorchInductor、临时文件和实验输出路由到 + 非系统盘 `${SPARSEVLLM_ARTIFACT_ROOT}`。 + +门禁： + +- 不覆盖现有 partial checkpoint。 +- checkpoint 不完整时不运行真实模型加载。 +- GPU 任务前重新检查全部设备，仅使用空闲设备。 + +### 阶段 1：Vendor 独立 MLA kernels + +首个提交： + +~~~ +feat: vendor lightllm mla kernels +~~~ + +内容： + +~~~ +src/sparsevllm/triton_kernel/mla/ + __init__.py + decode_stage1.py + decode_stage2.py + decode_schedule.py + copy_latent.py + gather_latent.py + README.md + LICENSE.lightllm +tests/test_mla_kernels.py +~~~ + +门禁： + +- 零 LightLLM runtime import。 +- 所有函数使用显式 tensor/stride/workspace/config API。 +- copy 正确跳过 slot=-1。 +- gather 支持 ragged、non-contiguous、完整历史和 padded rows。 +- Torch oracle 覆盖 local heads 5/10/20 和边界长度。 +- 空闲 H100 数值 smoke 通过后才合入。 + +该阶段不接 Model、CacheManager 或 OpenAI serving。 + +### 阶段 2：ComputeView tagged payload 迁移 + +工作： + +- 引入 AttentionViewMeta、ExplicitKVPayload、MlaLatentPayload。 +- 将现有 prefill/decode view 改为 meta + payload。 +- 迁移所有现有 manager/backend/tests 到 ExplicitKVPayload。 + +门禁： + +- 不加入 MLA 执行逻辑。 +- 当前 explicit-KV 测试全部保持通过。 +- git diff 中不存在 optional payload field bag。 +- 通用 backend 在 payload 错型时 fail fast。 + +### 阶段 3：Storage composition + +工作： + +- 引入 ExplicitKVStorage 和 MlaLatentStorage。 +- StandardCacheManager 委托 allocation、store、layer payload、bytes 和 accounting。 +- 现有 sparse managers 首版只接受 explicit layout。 + +门禁： + +- explicit storage 与迁移前分配量和结果一致。 +- MLA storage 容量按 576 BF16 values/token/layer 计算。 +- slot reuse/free/padding 不泄漏旧数据。 +- 所有不支持组合在配置阶段拒绝。 + +### 阶段 4：MLA operator/provider 与 decode + +工作： + +- 新增 MlaAttentionOpSpec、registry 和 Triton provider。 +- 初始化时绑定 provider。 +- 构造显式 workspace 和安全 launch config。 +- 打通 absorbed query -> latent decode -> V_b reconstruction。 + +门禁： + +- resolver supported/rejected tests。 +- 不在 forward 中 resolve 或动态选择 provider。 +- softmax scale 使用 256，而不是 576。 +- decode 与 Torch/Transformers component oracle 对齐。 +- static padded batch 不越界。 + +### 阶段 5：完整历史 chunked prefill + +工作： + +- 保存当前 chunk latent/rope。 +- gather 完整 visible history。 +- 临时展开 K/V。 +- 接现有 256 维 prefill attention kernel。 +- 释放 workset。 + +门禁： + +- 单 chunk 与多 chunk 输出对齐。 +- 第二及后续 chunk 能看到历史 token。 +- ragged batch、prefix length 和不同 chunk boundary 对齐。 +- workspace 有界且在不足时明确失败。 + +### 阶段 6：Tiny MLA/model contract + +工作： + +- 扩展 tiny-random config/checkpoint generator。 +- 建立两层模型：第 0 层 Dense、第 1 层 MoE。 +- 与 Transformers 比较 projection、RoPE、router、attention、MLP 和 logits。 + +门禁： + +- prefill logits 对齐。 +- 多步 decode greedy token 对齐。 +- TP rank-local projection slice 正确。 +- 所有 logits finite。 + +### 阶段 7：最小 packed-expert 抽取 + +工作： + +- 抽取 Qwen3-MoE packed expert 物理执行。 +- Qwen3-MoE 使用共享实现，保持语义不变。 +- GLM 准备复用相同 packed experts。 + +门禁： + +- Qwen3-MoE 现有 config、loader、provider、TP/EP 测试全部通过。 +- 不抽取 router、model block、checkpoint regex 或 MiniMax storage。 + +### 阶段 8：GLM MoE、模型与 loader + +工作： + +- 泛化 biased-sigmoid router 到 64/Top-4。 +- 实现 Dense/MoE/shared-expert topology。 +- 注册 model runner/config/runtime topology。 +- 实现 strict loader 和 MTP skip。 + +门禁： + +- tiny end-to-end 对齐。 +- loaded/skipped/missing tensor 统计满足精确约束。 +- Qwen3-MoE 和 MiniMax router 回归测试通过。 + +### 阶段 9：Transformers response template + +工作： + +- 增加 GLM declarative template fallback。 +- 加入非流式与流式 OpenAI 响应测试。 + +门禁： + +- 实际调用 Transformers parse_response / get_response_parser。 +- thinking、普通 content、多 tool calls、Unicode、跨 chunk tag、畸形输出均覆盖。 + +### 阶段 10：真实模型验证 + +顺序： + +1. checkpoint 完整性与 loader metadata 检查。 +2. HF reference 与 Sparse-vLLM 顺序执行，避免同时占用显存。 +3. TP1、eager、vanilla、短上下文 smoke。 +4. TP1 component/logit/token parity。 +5. 等待四张 GPU 同时空闲。 +6. TP4/EP1 smoke。 +7. TP1/TP4 greedy token 与 logits 对齐。 +8. OpenAI chat/reasoning/tool call smoke。 + +首轮真实模型建议： + +~~~ +max_model_len: 2048-4096 +prompt: 128-512 tokens +decode: 32-64 tokens +temperature: 0 +dtype: bfloat16 +decode_cuda_graph: false +enable_prefix_caching: false +vllm_sparse_method: "" +~~~ + +## 18. 测试矩阵 + +### 18.1 Kernel tests + +- batch：1、2、8。 +- local heads：5、10、20。 +- context：1、31、32、33、127、128、129、255、256、257、1024、4096。 +- slot：contiguous、ragged、non-contiguous、invalid duplicate rejection、padding -1。 +- dtype：BF16；其他 dtype 明确拒绝或单独验证。 +- stage1/stage2 分别与整体 oracle 对齐。 +- workspace 最小容量和不足容量失败。 + +### 18.2 ComputeView/storage tests + +- payload 类型不匹配失败。 +- explicit path 行为不变。 +- bytes-per-slot 与实际 tensor bytes 一致。 +- allocate/free/reuse 不泄漏。 +- accounting 覆盖所有 storage tensor。 +- 不支持的 sparse/graph/prefix 配置构造前失败。 + +### 18.3 Component parity + +- q_a/q_b projection。 +- kv_a norm 和 kv_b split。 +- 192+64 interleaved RoPE。 +- prefill expanded K/V。 +- decode K_b absorption 和 V_b reconstruction。 +- biased-sigmoid ids/weights/scaling。 +- routed + shared expert 输出。 +- Dense/MoE layer topology。 + +每个测试使用显式 atol/rtol。BF16 kernel 初始可用 3e-2 量级作为探测值，但最终容差必须由误差分布和独立 reference 决定，不能通过不断放宽容差让测试通过。 + +### 18.4 End-to-end parity + +- tiny prompt prefill logits。 +- tiny 多步 greedy decode。 +- 真实模型首 token logits/top-k margin。 +- 真实模型短序列 greedy token。 +- TP1 与 TP4 对齐。 + +非空输出不构成正确性证明。 + +## 19. 可观测性与实验产物 + +每次 GPU 验证至少保存： + +- code commit 和 dirty status。 +- 完整命令与 resolved config。 +- Python/Torch/Transformers/Triton/CUDA 版本。 +- GPU 型号和 device index。 +- model/checkpoint revision/hash/completeness。 +- selected MLA provider 与 rejection summary。 +- cache layout、每 token 字节数、实际 slot capacity。 +- loaded/skipped/missing weight counts。 +- raw outputs、parsed outputs、per-sample status、aggregate result。 +- stdout/stderr 和错误扫描。 + +建议通过环境变量配置非系统盘输出根目录： + +~~~ +${SPARSEVLLM_ARTIFACT_ROOT}/glm-4.7-flash// +~~~ + +系统盘空间不足时不得把 Hugging Face、Triton、TorchInductor 或实验临时文件写入默认 home cache。 + +## 20. 风险登记 + +| 严重度 | 风险 | 最小修正 | +| --- | --- | --- | +| P0 | chunked prefill 只展开当前 chunk，历史不可见 | 先 store，再按 active slots gather 完整历史，然后展开 | +| P0 | LightLLM copy 对 slot=-1 无 mask | vendor 时增加 padding mask 和回归测试 | +| P0 | 用 576 推导 softmax scale | OpSpec 独立记录 qk head dim 256 | +| P1 | latent 冒充 k_cache/v_cache | ComputeView tagged payload | +| P1 | model-specific CacheManager 形成组合爆炸 | CacheManager + storage composition | +| P1 | 原样复制 LightLLM wrapper | 只 vendor 数学核心，本地重写 provider/schedule | +| P1 | runtime kernel failure 后静默 fallback | 仅初始化时 resolver fallback；执行失败直接暴露 | +| P1 | 不支持组合落入 dense/default validation | GLM compatibility 在配置阶段显式拒绝 | +| P2 | prefill workspace 随历史无界增长 | 显式预算、容量估算和 fail-fast | +| P2 | H100 没有固定上游 tuning config | 安全静态 config + 实机数值/性能验证 | +| P2 | MoE 抽取范围过大 | 只抽 packed expert 物理执行 | +| P3 | vendor 来源/修改不清 | README、固定 commit、SPDX、license 和 source mapping | + +## 21. 预计文件影响 + +| 区域 | 文件 | +| --- | --- | +| ComputeView | src/sparsevllm/engine/cache_manager/base.py、所有 view 构造/消费点 | +| Storage | src/sparsevllm/engine/cache_manager/storage/*、standard.py | +| Attention layer | src/sparsevllm/layers/mla_attention.py、attention.py | +| Operator | src/sparsevllm/operators/mla_attention.py、registry tests | +| Kernels | src/sparsevllm/triton_kernel/mla/* | +| Model | src/sparsevllm/models/glm4_moe_lite.py | +| MoE | qwen3_moe.py、共享 packed-expert 模块、router kernel | +| Config/runtime | configs/model.py、configs/runtime.py、method_registry.py | +| Dispatch/warmup | engine/model_runner.py、engine/llm_engine.py | +| Cache factory | engine/cache_manager/base.py 或后续独立 factory | +| Serving | entrypoints/openai/serving/response_parsing.py、api_server.py | +| Tiny/tests | debug/tiny_random.py、新增 kernel/model/parser tests | + +## 22. 提交拆分 + +计划提交： + +1. feat: vendor lightllm mla kernels +2. refactor: tag attention payload views +3. feat: add mla cache storage +4. feat: add mla attention provider +5. feat: add mla prefill and decode +6. refactor: share packed moe experts +7. feat: add glm4 moe lite model +8. feat: add glm response template +9. test: validate glm4 runtime parity +10. docs: document glm4 support matrix + +每个提交要求： + +- git diff --check 通过。 +- 对应的 targeted tests 通过。 +- 不混入无关重构。 +- 不提交失败或不完整的真实实验结果作为成功结论。 + +ComputeView 迁移与 MLA 引入分开；Qwen MoE 抽取与 GLM model 引入分开；支持矩阵只在真实门禁完成后更新。 + +## 23. 完成定义 + +只有以下条件全部满足，才可标记“GLM-4.7-Flash 基础支持完成”： + +- Vendor kernel 有固定来源、license、修改说明和独立 oracle tests。 +- ComputeView 不再假定唯一 payload 是显式双 K/V。 +- 所有原有 explicit-KV 测试保持通过。 +- Persistent cache 只保存 512+64 latent/rope，不保存展开式 20-head K/V。 +- Chunked prefill 包含完整可见历史。 +- Padded decode 不对 slot=-1 写入。 +- MLA scale 使用 256 ** -0.5。 +- Provider 初始化时绑定，运行时无静默 fallback。 +- Tiny model 的 attention、MoE、logits 和多步 decode 与 Transformers 对齐。 +- 真实 checkpoint 除精确 MTP/远端 EP allowlist 外无 missing/unexpected weights。 +- TP1 真实短序列验证通过；TP4/EP1 在写入支持声明前必须另行验证。 +- OpenAI 非流式/流式 reasoning 和 tool-call parsing 通过 Transformers parser。 +- Artifact 能证明实际执行了目标 MLA provider、latent cache 和 GLM MoE 路径。 +- 文档支持矩阵准确区分已实现、已验证和暂不支持。 + +## 24. 开始执行前检查清单 + +- [ ] 当前分支同步到目标 main。 +- [ ] worktree 干净或已有改动已明确归属。 +- [ ] LightLLM source commit 固定。 +- [ ] vendor source/license manifest 准备完成。 +- [ ] GPU 状态重新检查。 +- [ ] 系统盘 cache/output 路径迁移到非系统盘 `${SPARSEVLLM_ARTIFACT_ROOT}`。 +- [ ] 首个 kernel commit 的 Torch oracle 写好。 +- [ ] ComputeView 迁移单独排期，不与 MLA 功能混合。 +- [ ] checkpoint 完整性门禁准备好。 +- [ ] 所有首版不支持组合有明确 config rejection。 + +第一步是独立完成并验证 feat: vendor lightllm mla kernels。该提交通过前，不进入 ComputeView、CacheManager 或模型实现。 + +## 25. 当前支持状态与验证摘要 + +当前里程碑已经完成 TP1、BF16、eager、vanilla latent MLA 基础支持。TP4/EP1 +未按当前范围执行，因此只能称为“代码路径已实现、尚未验证”，不能称为已支持。 + +### 25.1 已落地的架构边界 + +- MLA kernel 固定来源为 LightLLM commit + `65c174ee95ac6a6fd36b18b63d0b33d97e76b770`，本地保留 Apache-2.0 + license、逐文件来源映射和修改说明。 +- ComputeView 使用 tagged attention payload；explicit K/V 与 MLA latent payload + 通过类型契约区分，不使用互斥 optional tensor 或 metadata 字典旁路。 +- CacheManager 继续拥有 slot、生命周期和 logical view；storage strategy 分别管理 + explicit K/V 与每层每 token `512 + 64` BF16 latent/rope cache。 +- MLA provider 通过 operator resolver 在初始化时绑定；实机选择 + `triton_h100`，执行失败不会静默切换实现。 +- GLM MoE 复用 Qwen3-MoE 的 packed expert 物理执行，只保留 GLM 自己的 + biased-sigmoid router、Dense/MoE topology 和 checkpoint 语义。 +- Chat Completions 与 Responses API 均继续调用 Transformers response parser； + 本地只提供 GLM 声明式 template。Terminal EOS 和 stop boundary 在 dispatcher / + detokenizer 通用边界处理，不在 GLM parser 中写模型特判。 + +### 25.2 验证结果 + +| 门禁 | 结果 | 证据 | +| --- | --- | --- | +| CPU regression | `402 passed, 33 skipped, 87 subtests passed` | 完整 GLM/MLA/MoE/OpenAI/scheduler test selection | +| CUDA kernel/operator/model regression | `57 passed` | H100 kernel、operator、attention layer 与模型测试 | +| Tiny multi-chunk TP1 | 17-token prompt、chunk size 8，实际 3 chunks；prefill 与两步 decode argmax 全匹配，max abs diff 分别为 `0.00390625`、`0.00390625`、`0.0048828125` | run manifest、raw comparison 与 stdout | +| 真实 checkpoint TP1 | 48/48 shards、9491 tensors、55.77 GiB、8832 local expert shards、精确跳过 212 个 MTP tensors；prefill + 32 decode steps argmax 全匹配 | run manifest、raw comparison 与 stdout | +| OpenAI serving | 完整 thinking、thinking SSE 和真实 function tool call 通过；另以真实 checkpoint 复验 terminal EOS、单 token 内 stop boundary 及 Responses API | command/env manifest、raw token/text、parser-facing raw text、parsed responses 与 request logs | + +完整命令、依赖与设备版本、dirty status、raw/parsed output 和本机 artifact +路径保存在私有实验记录中，不写入公开仓库文档；上表只保留稳定的支持边界和 +可复现门禁摘要。 + +真实 checkpoint 的 32 步 greedy token 序列全部一致，足以证明该短序列的独立 +greedy rollout 一致。它不等价于严格逐元素 logits parity：后续步骤的 BF16 raw +logit max abs diff 可到 `6.0`，并带近似全局平移；当前证据只支持“argmax/greedy +sequence 一致”，不支持“所有 logits 数值严格等价”。 + +### 25.3 尚未完成的门禁 + +- TP4/EP1 真实 checkpoint smoke 与 TP1 对齐。 +- CUDA Graph、Prefix Cache/offload、量化、MTP 和所有 sparse method。 +- 128K/202K 长上下文容量与吞吐验证。 + +这些组合在完成独立验证前必须继续 fail fast 或维持未支持声明。 diff --git a/docs/zh/features/supported-models.md b/docs/zh/features/supported-models.md index 44331922..2d7fc40d 100644 --- a/docs/zh/features/supported-models.md +++ b/docs/zh/features/supported-models.md @@ -13,6 +13,7 @@ | Qwen3MoE | `qwen3_moe` | BF16 / FP16 / 块级 FP8 | ✅（TP > 1 时模型 dtype 仅支持 BF16） | 仅支持 1 | ✅ | | Qwen3.5 / Qwen3.6 | `qwen3_5` | BF16 / 块级 FP8 | ✅ | 仅支持 1 | 仅支持 1 | | Qwen3.6 MoE | `qwen3_5_moe` | BF16 / 块级 FP8 | ✅ | 仅支持 1 | ✅ | +| GLM-4.7-Flash | `glm4_moe_lite` | BF16 | 仅支持 1 | 仅支持 1 | 仅支持 1 | | Llama 3 / 3.1 | `llama` | BF16 / FP16 | ✅ | 仅支持 1 | 仅支持 1 | | MiniMax M2.7 | `minimax_m2` | 块级 FP8，非量化权重使用 BF16 | ✅ | 仅支持 1 | ✅ | @@ -29,6 +30,11 @@ TP 要求模型 dtype 为 BF16；FP16 Qwen3MoE checkpoint 仅支持 `TP=1`。当 Qwen3.5/Qwen3.6 Dense 配置在内部统一规范为 `model_type=qwen3_5`；Qwen3.6 MoE 使用 `model_type=qwen3_5_moe`。 +GLM-4.7-Flash 当前仅支持 BF16、eager、vanilla latent MLA 和 +`TP=DP=EP=1`。loader 会有意跳过 checkpoint 中的 MTP 层。TP 大于 1、 +expert parallel、CUDA Graph、Prefix Cache、量化和稀疏方法必须经过独立验证后 +才会进入支持范围。 + ## 稀疏方法支持 | 模型 | Vanilla | StreamingLLM | SnapKV | H2O | PyramidKV | OmniKV | QuEST | R-KV | SkipKV | DeltaKV | @@ -38,6 +44,7 @@ MoE 使用 `model_type=qwen3_5_moe`。 | Qwen3MoE | ✅ | ✅ | ✅ | 实验性⁴ | ✅ | ✅ | ✅ | ✅ | — | — | | Qwen3.5 / Qwen3.6 | ✅ | ✅ | ✅ | 实验性⁴ | ✅ | ✅ | ✅ | ✅ | — | 匹配的 checkpoint³ | | Qwen3.6 MoE | ✅ | ✅ | ✅ | 实验性⁴ | ✅ | ✅ | ✅ | ✅ | — | — | +| GLM-4.7-Flash | ✅ | — | — | — | — | — | — | — | — | — | | Llama 3 / 3.1 | ✅ | ✅ | ✅ | 实验性⁴ | ✅ | ✅ | ✅ | ✅ | 指定 checkpoint¹ | 需要 compressor² | | MiniMax M2.7 | ✅ | ✅ | ✅ | 实验性⁴ | ✅ | ✅ | ✅ | ✅ | — | — | From 34639aa75459ce93bbc5831161023fd7ca5c0b22 Mon Sep 17 00:00:00 2001 From: Codex Date: Wed, 5 Aug 2026 03:56:40 +0800 Subject: [PATCH 15/48] docs: record glm tp validation --- docs/en/features/supported-models.md | 7 ++-- docs/zh/design/glm-4.7-flash-support-plan.md | 36 +++++++++++++------- docs/zh/features/supported-models.md | 10 +++--- 3 files changed, 32 insertions(+), 21 deletions(-) diff --git a/docs/en/features/supported-models.md b/docs/en/features/supported-models.md index c2b50f7f..65788683 100644 --- a/docs/en/features/supported-models.md +++ b/docs/en/features/supported-models.md @@ -17,7 +17,7 @@ parallel size must use that value. | Qwen3MoE | `qwen3_moe` | BF16 / FP16 / block FP8 | ✅ (TP > 1: BF16 model dtype only) | 1 only | ✅ | | Qwen3.5 / Qwen3.6 | `qwen3_5` | BF16 / block FP8 | ✅ | 1 only | 1 only | | Qwen3.6 MoE | `qwen3_5_moe` | BF16 / block FP8 | ✅ | 1 only | ✅ | -| GLM-4.7-Flash | `glm4_moe_lite` | BF16 | 1 only | 1 only | 1 only | +| GLM-4.7-Flash | `glm4_moe_lite` | BF16 | 1 / 2 / 4 (H100 only) | 1 only | 1 only | | Llama 3 / 3.1 | `llama` | BF16 / FP16 | ✅ | 1 only | 1 only | | MiniMax M2.7 | `minimax_m2` | block FP8 with BF16 non-quantized weights | ✅ | 1 only | ✅ | @@ -39,8 +39,9 @@ normalized internally to `model_type=qwen3_5`; Qwen3.6 MoE uses `model_type=qwen3_5_moe`. GLM-4.7-Flash support is currently limited to BF16, eager execution, vanilla -latent MLA, and `TP=DP=EP=1`. The loader intentionally skips the checkpoint's -MTP layer. TP greater than 1, expert parallelism, CUDA Graph, prefix caching, +latent MLA, `TP=1/2/4`, `DP=1`, and `EP=1` on NVIDIA H100 80GB HBM3. The +loader intentionally skips the checkpoint's MTP layer. Other GPU +architectures, TP sizes, expert parallelism, CUDA Graph, prefix caching, quantization, and sparse methods remain unsupported until separately validated. ## Sparse Method Support diff --git a/docs/zh/design/glm-4.7-flash-support-plan.md b/docs/zh/design/glm-4.7-flash-support-plan.md index f1b50b62..a5192398 100644 --- a/docs/zh/design/glm-4.7-flash-support-plan.md +++ b/docs/zh/design/glm-4.7-flash-support-plan.md @@ -1,7 +1,7 @@ # GLM-4.7-Flash 支持实施计划 -- 状态：Implemented（TP1 已验证；TP4/EP1 待验证） -- 更新时间：2026-08-04 +- 状态：Implemented（H100 TP1/TP2/TP4、EP1 已验证） +- 更新时间：2026-08-05 - 目标分支：codex/glm-4.7-flash - 目标模型：zai-org/GLM-4.7-Flash / glm4_moe_lite @@ -527,7 +527,8 @@ TP 布局建议： - 每 rank 持有本 rank 的 K_b/V_b head slices。 - persistent latent/rope cache 在各 TP rank 上语义一致。 -首版验证 TP1 和 TP4；TP2 可以保留配置能力，但未经实测不列为已验证。 +首版门禁以 TP1 和 TP4 为目标；本轮同时补测 TP2。支持声明只列入真实 +checkpoint 已实测的 `TP=1/2/4`，不从配置可构造性外推其他 TP 规模。 ## 13. MoE 复用边界 @@ -626,10 +627,12 @@ CLI 建议新增通用 auto；现有 qwen3、minimax_m2 保留兼容，glm47 仅 | 能力 | 首版状态 | 门禁 | | --- | --- | --- | | BF16 | 已实现并验证 | tiny + 真实权重验证 | +| NVIDIA H100 80GB HBM3 | 已实现并验证 | `triton_h100` provider 实机门禁 | | eager | 已实现并验证 | 必须显式启用 | | vanilla latent MLA | 已实现并验证 | kernel/storage/model 全链路 | | TP1 | 已实现并验证 | 真实模型 | -| TP4/EP1 | 已实现路径，尚未验证 | 验证前不写入支持声明 | +| TP2/EP1 | 已实现并验证 | 真实模型 32 步 greedy 对齐 | +| TP4/EP1 | 已实现并验证 | 真实模型、multi-chunk、ragged batch 与 API 并发 | | EP > 1 | 暂不支持 | 后续独立验证 | | CUDA Graph | 暂不支持 | 配置阶段拒绝 | | Prefix Cache/offload | 暂不支持 | 配置阶段拒绝 | @@ -977,7 +980,7 @@ ComputeView 迁移与 MLA 引入分开；Qwen MoE 抽取与 GLM model 引入分 - Provider 初始化时绑定，运行时无静默 fallback。 - Tiny model 的 attention、MoE、logits 和多步 decode 与 Transformers 对齐。 - 真实 checkpoint 除精确 MTP/远端 EP allowlist 外无 missing/unexpected weights。 -- TP1 真实短序列验证通过；TP4/EP1 在写入支持声明前必须另行验证。 +- TP1/TP2/TP4、EP1 真实短序列验证通过；其他 TP/EP 规模不得外推。 - OpenAI 非流式/流式 reasoning 和 tool-call parsing 通过 Transformers parser。 - Artifact 能证明实际执行了目标 MLA provider、latent cache 和 GLM MoE 路径。 - 文档支持矩阵准确区分已实现、已验证和暂不支持。 @@ -999,8 +1002,9 @@ ComputeView 迁移与 MLA 引入分开；Qwen MoE 抽取与 GLM model 引入分 ## 25. 当前支持状态与验证摘要 -当前里程碑已经完成 TP1、BF16、eager、vanilla latent MLA 基础支持。TP4/EP1 -未按当前范围执行，因此只能称为“代码路径已实现、尚未验证”，不能称为已支持。 +当前里程碑已经完成 BF16、eager、vanilla latent MLA，以及 +`TP=1/2/4`、`DP=1`、`EP=1` 的基础支持，硬件限定为 NVIDIA H100 80GB +HBM3。支持声明只覆盖这些实测组合。 ### 25.1 已落地的架构边界 @@ -1027,21 +1031,27 @@ ComputeView 迁移与 MLA 引入分开；Qwen MoE 抽取与 GLM model 引入分 | CUDA kernel/operator/model regression | `57 passed` | H100 kernel、operator、attention layer 与模型测试 | | Tiny multi-chunk TP1 | 17-token prompt、chunk size 8，实际 3 chunks；prefill 与两步 decode argmax 全匹配，max abs diff 分别为 `0.00390625`、`0.00390625`、`0.0048828125` | run manifest、raw comparison 与 stdout | | 真实 checkpoint TP1 | 48/48 shards、9491 tensors、55.77 GiB、8832 local expert shards、精确跳过 212 个 MTP tensors；prefill + 32 decode steps argmax 全匹配 | run manifest、raw comparison 与 stdout | -| OpenAI serving | 完整 thinking、thinking SSE 和真实 function tool call 通过；另以真实 checkpoint 复验 terminal EOS、单 token 内 stop boundary 及 Responses API | command/env manifest、raw token/text、parser-facing raw text、parsed responses 与 request logs | +| 真实 checkpoint TP2 | 每 rank 28.01 GiB，两个 rank 均加载 48/48 shards、9491 tensors；prefill + 32 decode steps 共 33/33 argmax 全匹配，greedy token 序列与 TP1 相同 | run manifest、raw comparison 与 stdout | +| 真实 checkpoint TP4/EP1 | 每 rank 14.13 GiB，四个 rank 均绑定 `triton_h100` MLA 与 Triton MoE；短序列 33/33、257-token 五段 prefill + decode 9/9 argmax 全匹配 | run manifest、raw comparison 与 stdout | +| TP4 synthetic API ragged + natural LCC | 64/256/768/1536-token synthetic API batch 连续三轮共 12/12 成功且跨轮逐字稳定；另以 1004/1023/1024/1025-token natural LCC 样本对 HF teacher-forced prefill 7/7、decode 27/28，唯一 teacher-forced argmax mismatch 为两个 HF 候选 logit 完全相等的 BF16 tie | raw/parsed responses、per-row comparison 与 request logs | +| TP4 OpenAI client gates | Chat 普通/思考、Chat SSE、Responses、Responses SSE、function tool call、streaming tool call 和 4-way concurrency 共 8/8 client gates 通过；parser 实际委托 Transformers。服务在验证后由 operator 主动停止，teardown 不作为 client gate | command/env manifest、raw parser events、parsed responses、validation status 与 request logs | 完整命令、依赖与设备版本、dirty status、raw/parsed output 和本机 artifact 路径保存在私有实验记录中，不写入公开仓库文档；上表只保留稳定的支持边界和 可复现门禁摘要。 -真实 checkpoint 的 32 步 greedy token 序列全部一致，足以证明该短序列的独立 -greedy rollout 一致。它不等价于严格逐元素 logits parity：后续步骤的 BF16 raw -logit max abs diff 可到 `6.0`，并带近似全局平移；当前证据只支持“argmax/greedy -sequence 一致”，不支持“所有 logits 数值严格等价”。 +TP1、TP2 和 TP4 的真实 checkpoint 32 步 greedy token 序列全部一致，足以证明 +该短序列的独立 greedy rollout 一致。它不等价于严格逐元素 logits parity： +后续步骤及 multi-chunk case 的 BF16 raw logit max abs diff 可到 `9.0`，并带 +近似全局平移；自然文本 teacher-forced 测试还观测到一次零 margin tie。当前证据 +只支持“非 tie 的 argmax 与短序列 greedy 行为一致”，不支持“所有 logits 数值 +严格等价”。 ### 25.3 尚未完成的门禁 -- TP4/EP1 真实 checkpoint smoke 与 TP1 对齐。 - CUDA Graph、Prefix Cache/offload、量化、MTP 和所有 sparse method。 - 128K/202K 长上下文容量与吞吐验证。 +- `TP=5` 等其他可整除配置，以及 `EP > 1`。 +- H100 以外的 GPU 架构。 这些组合在完成独立验证前必须继续 fail fast 或维持未支持声明。 diff --git a/docs/zh/features/supported-models.md b/docs/zh/features/supported-models.md index 2d7fc40d..21819784 100644 --- a/docs/zh/features/supported-models.md +++ b/docs/zh/features/supported-models.md @@ -13,7 +13,7 @@ | Qwen3MoE | `qwen3_moe` | BF16 / FP16 / 块级 FP8 | ✅（TP > 1 时模型 dtype 仅支持 BF16） | 仅支持 1 | ✅ | | Qwen3.5 / Qwen3.6 | `qwen3_5` | BF16 / 块级 FP8 | ✅ | 仅支持 1 | 仅支持 1 | | Qwen3.6 MoE | `qwen3_5_moe` | BF16 / 块级 FP8 | ✅ | 仅支持 1 | ✅ | -| GLM-4.7-Flash | `glm4_moe_lite` | BF16 | 仅支持 1 | 仅支持 1 | 仅支持 1 | +| GLM-4.7-Flash | `glm4_moe_lite` | BF16 | 1 / 2 / 4（仅 H100） | 仅支持 1 | 仅支持 1 | | Llama 3 / 3.1 | `llama` | BF16 / FP16 | ✅ | 仅支持 1 | 仅支持 1 | | MiniMax M2.7 | `minimax_m2` | 块级 FP8，非量化权重使用 BF16 | ✅ | 仅支持 1 | ✅ | @@ -30,10 +30,10 @@ TP 要求模型 dtype 为 BF16；FP16 Qwen3MoE checkpoint 仅支持 `TP=1`。当 Qwen3.5/Qwen3.6 Dense 配置在内部统一规范为 `model_type=qwen3_5`；Qwen3.6 MoE 使用 `model_type=qwen3_5_moe`。 -GLM-4.7-Flash 当前仅支持 BF16、eager、vanilla latent MLA 和 -`TP=DP=EP=1`。loader 会有意跳过 checkpoint 中的 MTP 层。TP 大于 1、 -expert parallel、CUDA Graph、Prefix Cache、量化和稀疏方法必须经过独立验证后 -才会进入支持范围。 +GLM-4.7-Flash 当前仅支持 BF16、eager、vanilla latent MLA、`TP=1/2/4`、 +`DP=1` 和 `EP=1`，硬件限定为 NVIDIA H100 80GB HBM3。loader 会有意跳过 +checkpoint 中的 MTP 层。其他 GPU 架构、TP 规模、expert parallel、CUDA Graph、 +Prefix Cache、量化和稀疏方法必须经过独立验证后才会进入支持范围。 ## 稀疏方法支持 From a2213c8ffbc94f720b69cdaa0e658c1ceb6801a3 Mon Sep 17 00:00:00 2001 From: Codex Date: Wed, 5 Aug 2026 16:38:04 +0800 Subject: [PATCH 16/48] feat: expand glm runtime support --- benchmark/long_bench/pred.py | 42 +- docs/en/features/supported-models.md | 36 +- docs/zh/design/glm-4.7-flash-support-plan.md | 77 +- docs/zh/features/supported-models.md | 33 +- scripts/benchmarks/bench_prefix_cache.py | 117 ++ .../compare_decode_graph_eager_logits.py | 823 ++++++++-- .../compare_glm_parallel_variants.py | 293 ++++ .../validate_glm_checkpoint_parallel.py | 546 +++++++ src/sparsevllm/configs/model.py | 59 +- src/sparsevllm/configs/prefix_cache.py | 4 +- .../engine/cache_manager/__init__.py | 2 + src/sparsevllm/engine/cache_manager/base.py | 152 +- src/sparsevllm/engine/cache_manager/h2o.py | 99 +- src/sparsevllm/engine/cache_manager/rkv.py | 133 +- src/sparsevllm/engine/cache_manager/snapkv.py | 112 +- .../engine/cache_manager/standard.py | 2 +- .../engine/cache_manager/storage/base.py | 14 + .../cache_manager/storage/explicit_kv.py | 52 + .../cache_manager/storage/mla_latent.py | 104 +- src/sparsevllm/engine/chain_cache.py | 6 +- src/sparsevllm/engine/decode_cuda_graph.py | 18 + src/sparsevllm/engine/model_runner.py | 129 +- src/sparsevllm/engine/sparse_controller.py | 113 +- src/sparsevllm/layers/mla_attention.py | 110 +- src/sparsevllm/method_registry.py | 24 +- src/sparsevllm/models/glm4_moe_lite.py | 109 +- src/sparsevllm/models/spec.py | 2 + src/sparsevllm/operators/mla_attention.py | 59 +- .../triton_kernel/mla/decode_schedule.py | 27 +- .../triton_kernel/mla/decode_stage1.py | 114 +- src/sparsevllm/triton_kernel/omnikv_fused.py | 55 +- src/sparsevllm/utils/context.py | 1 + tests/glm_test_helpers.py | 86 ++ tests/test_attention_cache_storage.py | 308 ++-- tests/test_chain_prefix_cache.py | 1 + .../test_compare_decode_graph_eager_logits.py | 246 +++ tests/test_compare_glm_parallel_variants.py | 104 ++ tests/test_glm4_moe_lite.py | 291 +++- tests/test_glm_cuda_graph.py | 1363 +++++++++++++++++ tests/test_glm_mla_prefix_cache.py | 517 +++++++ tests/test_glm_mla_sparse_methods.py | 343 +++++ tests/test_glm_runtime_compatibility.py | 297 ++++ tests/test_h2o_cache_manager.py | 40 + tests/test_longbench_deltakv_contracts.py | 59 + tests/test_mla_attention_layer.py | 93 +- tests/test_mla_attention_operator.py | 85 +- tests/test_mla_kernels.py | 206 +++ tests/test_omnikv_decode_score_lifecycle.py | 48 +- tests/test_prefix_cache.py | 189 ++- tests/test_prefix_cache_bench.py | 94 ++ tests/test_rkv_skipkv_methods.py | 352 +++++ tests/test_sparse_state_summary.py | 114 +- tests/test_sparsevllm_regression_grading.py | 25 +- tests/test_sparsevllm_regression_suite.py | 46 +- .../test_validate_glm_checkpoint_parallel.py | 198 +++ 55 files changed, 8122 insertions(+), 450 deletions(-) create mode 100644 scripts/validation/compare_glm_parallel_variants.py create mode 100644 scripts/validation/validate_glm_checkpoint_parallel.py create mode 100644 tests/glm_test_helpers.py create mode 100644 tests/test_compare_decode_graph_eager_logits.py create mode 100644 tests/test_compare_glm_parallel_variants.py create mode 100644 tests/test_glm_cuda_graph.py create mode 100644 tests/test_glm_mla_prefix_cache.py create mode 100644 tests/test_glm_mla_sparse_methods.py create mode 100644 tests/test_glm_runtime_compatibility.py create mode 100644 tests/test_validate_glm_checkpoint_parallel.py diff --git a/benchmark/long_bench/pred.py b/benchmark/long_bench/pred.py index 59fc46fe..27e94a1a 100644 --- a/benchmark/long_bench/pred.py +++ b/benchmark/long_bench/pred.py @@ -27,6 +27,7 @@ BASE_PATH = os.getenv("SPARSEVLLM_OUTPUT_DIR", str(REPO_ROOT / "outputs")) DATA_PREFIX_PATH = os.getenv("SPARSEVLLM_LONGBENCH_DATA_DIR") or os.getenv("SPARSEVLLM_DATA_DIR") +DEFAULT_MAX_MODEL_LEN = 121_000 NO_CHAT_TEMPLATE_DATASETS = {"trec", "triviaqa", "samsum", "lsht", "lcc", "repobench-p"} SAMPLE_STATUSES = { "success", @@ -148,10 +149,9 @@ def _artifact_paths(out_root: str) -> dict[str, str]: } -def _write_decode_cuda_graph_status( +def _decode_cuda_graph_status( *, generate_fn, - out_root: str, rank: int, ) -> dict[str, Any]: llm = getattr(generate_fn, "_sparsevllm_llm", None) @@ -181,9 +181,35 @@ def _write_decode_cuda_graph_status( "state_count": int(len(graph_states)), "graph_count": int(graph_count), "active": bool(graph_count > 0), + "capture_count": int(getattr(graph_runner, "capture_count", 0)), + "replay_count": int(getattr(graph_runner, "replay_count", 0)), + "eager_static_count": int(getattr(graph_runner, "eager_static_count", 0)), + "force_eager_count": int(getattr(graph_runner, "force_eager_count", 0)), "last_state_key": str(getattr(graph_runner, "last_state_key", None)), "state_keys": [str(key) for key in graph_states], } + return graph_status + + +def _write_decode_cuda_graph_status( + *, + generate_fn, + out_root: str, + rank: int, + before: dict[str, Any] | None = None, +) -> dict[str, Any]: + graph_status = _decode_cuda_graph_status(generate_fn=generate_fn, rank=rank) + if before is not None: + counter_keys = ( + "capture_count", + "replay_count", + "eager_static_count", + "force_eager_count", + ) + graph_status["before"] = before + graph_status["counter_delta"] = { + key: int(graph_status[key]) - int(before[key]) for key in counter_keys + } status_path = os.path.join( out_root, f"decode_cuda_graph_status_rank{rank}.json", @@ -493,6 +519,7 @@ def get_pred(rank, data, dataset_info, args, model, tokenizer, model_max_length, def worker(rank, world_size, datasets, dataset2prompt, dataset2maxlen, args, out_root, max_length_limit): seed_everything(42) model, tokenizer, model_max_length, eos_token_ids = load_model_and_tokenizer(rank, args) + graph_status_before = _decode_cuda_graph_status(generate_fn=model, rank=rank) for dataset in datasets: data_path = get_longbench_data_path(dataset, args.e) @@ -584,6 +611,7 @@ def worker(rank, world_size, datasets, dataset2prompt, dataset2maxlen, args, out generate_fn=model, out_root=out_root, rank=rank, + before=graph_status_before, ) @@ -668,6 +696,12 @@ def parse_args(): parser.add_argument("--worker_rank", type=int, default=-1) parser.add_argument("--worker_world_size", type=int, default=1) parser.add_argument("--output_root", type=str, default=None) + parser.add_argument( + "--max_model_len", + type=int, + default=None, + help="Runtime context limit (default: 121000).", + ) return parser.parse_args() @@ -711,7 +745,9 @@ def parse_args(): with open(os.path.join(out_root, artifact), "w", encoding="utf-8") as f: pass - max_length_limit = 120_000 + 1000 + max_length_limit = DEFAULT_MAX_MODEL_LEN if args.max_model_len is None else args.max_model_len + if max_length_limit <= 0: + raise ValueError(f"--max_model_len must be > 0, got {max_length_limit}.") args.max_model_len = max_length_limit if args.worker_rank < 0: diff --git a/docs/en/features/supported-models.md b/docs/en/features/supported-models.md index 65788683..fd374fe3 100644 --- a/docs/en/features/supported-models.md +++ b/docs/en/features/supported-models.md @@ -17,7 +17,7 @@ parallel size must use that value. | Qwen3MoE | `qwen3_moe` | BF16 / FP16 / block FP8 | ✅ (TP > 1: BF16 model dtype only) | 1 only | ✅ | | Qwen3.5 / Qwen3.6 | `qwen3_5` | BF16 / block FP8 | ✅ | 1 only | 1 only | | Qwen3.6 MoE | `qwen3_5_moe` | BF16 / block FP8 | ✅ | 1 only | ✅ | -| GLM-4.7-Flash | `glm4_moe_lite` | BF16 | 1 / 2 / 4 (H100 only) | 1 only | 1 only | +| GLM-4.7-Flash | `glm4_moe_lite` | BF16 | 1 / 2 / 4 (H100 only)⁵ | 1 only | 1 / 2 / 4⁵ | | Llama 3 / 3.1 | `llama` | BF16 / FP16 | ✅ | 1 only | 1 only | | MiniMax M2.7 | `minimax_m2` | block FP8 with BF16 non-quantized weights | ✅ | 1 only | ✅ | @@ -26,23 +26,28 @@ including the attention heads and vocabulary size, to be divisible by the selected TP size. Qwen3MoE and MiniMax M2.7 use a hybrid layout when the outer TP size `T` is -greater than 1: attention TP is `T`, MoE EP is `E`, MoE TP is `T / E`, and -the distributed world size is `T`. This layout requires `DP=1` and `T % E == -0`. The expert count must be divisible by `E`, and the MoE intermediate -dimension must be divisible by `T / E`. Qwen3MoE outer TP requires a BF16 -model dtype; FP16 Qwen3MoE checkpoints are limited to `TP=1`. When `TP=1`, -the existing EP layout uses world size `E`. +greater than 1. GLM-4.7-Flash uses the same layout when both `T > 1` and +`E > 1`: attention TP is `T`, MoE EP is `E`, MoE TP is `T / E`, and the +distributed world size is `T`. This layout requires `DP=1` and `T % E == 0`. +The expert count must be divisible by `E`, and the MoE intermediate dimension +must be divisible by `T / E`. Qwen3MoE outer TP requires a BF16 model dtype; +FP16 Qwen3MoE checkpoints are limited to `TP=1`. When `TP=1`, the existing EP +layout uses world size `E`. Block FP8 support requires E4M3 weights, dynamic activation quantization, and a `128 x 128` weight block size. Qwen3.5/Qwen3.6 dense configurations are normalized internally to `model_type=qwen3_5`; Qwen3.6 MoE uses `model_type=qwen3_5_moe`. -GLM-4.7-Flash support is currently limited to BF16, eager execution, vanilla -latent MLA, `TP=1/2/4`, `DP=1`, and `EP=1` on NVIDIA H100 80GB HBM3. The -loader intentionally skips the checkpoint's MTP layer. Other GPU -architectures, TP sizes, expert parallelism, CUDA Graph, prefix caching, -quantization, and sparse methods remain unsupported until separately validated. +GLM-4.7-Flash uses BF16 latent MLA on NVIDIA H100 80GB HBM3 and requires +`DP=1` plus `enforce_eager=True`. The validated `(TP, EP)` layouts are +`(1,1)`, `(2,1)`, `(4,1)`, `(1,2)`, `(1,4)`, `(2,2)`, `(4,2)`, and `(4,4)`. +Across all eight layouts, vanilla, StreamingLLM, SnapKV, H2O, OmniKV, and R-KV +support decode CUDA Graph and prefix caching together. Prefix caching uses +radix mode for vanilla and OmniKV, and chain mode for StreamingLLM, SnapKV, +H2O, and R-KV. Prefix offload, quantization, and the other sparse methods +remain unsupported. The loader intentionally skips the checkpoint's MTP +layer. ## Sparse Method Support @@ -53,7 +58,7 @@ quantization, and sparse methods remain unsupported until separately validated. | Qwen3MoE | ✅ | ✅ | ✅ | Experimental⁴ | ✅ | ✅ | ✅ | ✅ | — | — | | Qwen3.5 / Qwen3.6 | ✅ | ✅ | ✅ | Experimental⁴ | ✅ | ✅ | ✅ | ✅ | — | Matched checkpoint³ | | Qwen3.6 MoE | ✅ | ✅ | ✅ | Experimental⁴ | ✅ | ✅ | ✅ | ✅ | — | — | -| GLM-4.7-Flash | ✅ | — | — | — | — | — | — | — | — | — | +| GLM-4.7-Flash | ✅⁵ | ✅⁵ | ✅⁵ | Experimental⁴⁵ | — | ✅⁵ | — | ✅⁵ | — | — | | Llama 3 / 3.1 | ✅ | ✅ | ✅ | Experimental⁴ | ✅ | ✅ | ✅ | ✅ | Selected checkpoint¹ | Compressor required² | | MiniMax M2.7 | ✅ | ✅ | ✅ | Experimental⁴ | ✅ | ✅ | ✅ | ✅ | — | — | @@ -71,4 +76,9 @@ scores and retains tokens using its local heads or KV heads, without cross-rank sparse-index aggregation. This is not guaranteed to be equivalent to TP=1 or global-head selection. Model-specific TP, EP, and DP restrictions still apply. +⁵ GLM support is limited to the eight `(TP, EP)` layouts listed above with +`DP=1`. At `TP>1`, head-scored sparse methods use TP-local selection without +cross-rank sparse-index aggregation, so their selection semantics are not +guaranteed to match `TP=1`. + `—` means that the combination is not currently supported. diff --git a/docs/zh/design/glm-4.7-flash-support-plan.md b/docs/zh/design/glm-4.7-flash-support-plan.md index a5192398..9132c589 100644 --- a/docs/zh/design/glm-4.7-flash-support-plan.md +++ b/docs/zh/design/glm-4.7-flash-support-plan.md @@ -1,6 +1,6 @@ # GLM-4.7-Flash 支持实施计划 -- 状态：Implemented（H100 TP1/TP2/TP4、EP1 已验证） +- 状态：Implemented（H100 BF16；8 种 TP/EP 布局 × 6 种方法的 Graph + Prefix 组合已验证） - 更新时间：2026-08-05 - 目标分支：codex/glm-4.7-flash - 目标模型：zai-org/GLM-4.7-Flash / glm4_moe_lite @@ -19,7 +19,12 @@ 6. MoE 只抽取 Qwen3-MoE 中真正通用的 packed-expert 物理执行部分，不构建万能 MoE 模型框架。 7. GLM response parsing 继续委托 Transformers；本地只提供缺失的声明式 response template。 -首版目标是 BF16、eager、vanilla latent MLA、TP1 和 TP4/EP1 的短上下文正确性。CUDA Graph、Prefix Cache、稀疏方法、MTP、量化和实用长上下文能力均需后续单独验证。 +本计划最初以 BF16、eager、vanilla latent MLA、TP1 和 TP4/EP1 的短上下文 +正确性为首版目标。后续扩展已在相同 H100/BF16 边界内验证 8 种合法 +`(TP, EP)` 布局，以及 vanilla、StreamingLLM、SnapKV、H2O、OmniKV、 +R-KV 与 decode CUDA Graph、Prefix Cache 的完整交叉组合。当前精确支持边界 +以第 16 节和第 25 节为准；MTP、量化、prefix offload 与实用长上下文仍需 +独立验证。 ## 2. 目标与非目标 @@ -48,6 +53,9 @@ - LightLLM Python package 或 runtime 依赖。 - 在 kernel 失败后切换到 Torch/其他 backend 的运行时 fallback。 +以上列表记录首版冻结范围，不代表后续扩展仍全部未支持。已完成的后续能力和 +仍然保留的拒绝边界见第 16 节兼容矩阵。 + ## 3. 已知模型契约 官方配置来源： @@ -622,21 +630,25 @@ tokenizer.get_response_parser() # 流式 CLI 建议新增通用 auto；现有 qwen3、minimax_m2 保留兼容，glm47 仅作为 alias，不选择另一套 parser 实现。 -## 16. 首版兼容矩阵 +## 16. 当前兼容矩阵 -| 能力 | 首版状态 | 门禁 | +| 能力 | 当前状态 | 门禁 | | --- | --- | --- | | BF16 | 已实现并验证 | tiny + 真实权重验证 | | NVIDIA H100 80GB HBM3 | 已实现并验证 | `triton_h100` provider 实机门禁 | -| eager | 已实现并验证 | 必须显式启用 | +| eager | 已实现并验证 | `enforce_eager=True` | | vanilla latent MLA | 已实现并验证 | kernel/storage/model 全链路 | -| TP1 | 已实现并验证 | 真实模型 | +| TP1/EP1 | 已实现并验证 | 真实模型 | | TP2/EP1 | 已实现并验证 | 真实模型 32 步 greedy 对齐 | | TP4/EP1 | 已实现并验证 | 真实模型、multi-chunk、ragged batch 与 API 并发 | -| EP > 1 | 暂不支持 | 后续独立验证 | -| CUDA Graph | 暂不支持 | 配置阶段拒绝 | -| Prefix Cache/offload | 暂不支持 | 配置阶段拒绝 | -| sparse methods | 暂不支持 | method registry 显式拒绝 | +| TP1/EP2、TP1/EP4 | 已实现并验证 | replicated attention + EP，跨 rank router/MoE/logits 一致性与 expert 覆盖 | +| TP2/EP2、TP4/EP2、TP4/EP4 | 已实现并验证 | outer TP：attention TP=`T`、MoE EP=`E`、MoE TP=`T/E`、world=`T` | +| decode CUDA Graph | 8 种布局均支持 6 种方法 | 每个 world rank 的 capture/replay、零 eager/force-eager fallback | +| Prefix Cache | 8 种布局均支持 6 种方法 | vanilla/OmniKV radix；StreamingLLM/SnapKV/H2O/R-KV chain | +| CUDA Graph + Prefix Cache | 8 种布局 × 6 种方法已联合验证 | 48/48 真实 checkpoint 组合门禁 | +| sparse methods | 8 种布局均支持 StreamingLLM、SnapKV、H2O、OmniKV、R-KV | H2O 维持实验性标记；TP>1 使用 TP-local selection | +| Prefix offload | 暂不支持 | 配置阶段拒绝 | +| PyramidKV、QuEST、SkipKV、DeltaKV | 暂不支持 | method registry 显式拒绝 | | MTP | 暂不支持 | loader 精确跳过 | | quantization | 暂不支持 | checkpoint config 拒绝 | | 128K/202K | 未验证 | 不写入支持声明 | @@ -882,6 +894,15 @@ vllm_sparse_method: "" - 真实模型首 token logits/top-k margin。 - 真实模型短序列 greedy token。 - TP1 与 TP4 对齐。 +- TP1/EP2 与 TP1/EP4 的跨 rank router、MoE 输出和 logits 一致性。 +- 8 种合法 `(TP, EP)` 布局的 group topology、expert 分片与执行覆盖。 +- vanilla/OmniKV radix 与 StreamingLLM/SnapKV/H2O/R-KV chain prefix cache + 的真实 hit 和 token reuse。 +- vanilla、StreamingLLM、SnapKV、H2O、OmniKV、R-KV 的真实模型 + static/eager 与 decode CUDA Graph 全词表 logits 对比；同时要求 Graph capture、 + business-request replay、零 eager/force-eager fallback 和方法特定触发证据。 +- 以上 6 种方法与 8 种并行布局、decode CUDA Graph、Prefix Cache 的完整 + 交叉组合；不能用单特性分别通过替代联合生命周期门禁。 非空输出不构成正确性证明。 @@ -1002,9 +1023,12 @@ ComputeView 迁移与 MLA 引入分开；Qwen MoE 抽取与 GLM model 引入分 ## 25. 当前支持状态与验证摘要 -当前里程碑已经完成 BF16、eager、vanilla latent MLA，以及 -`TP=1/2/4`、`DP=1`、`EP=1` 的基础支持，硬件限定为 NVIDIA H100 80GB -HBM3。支持声明只覆盖这些实测组合。 +当前里程碑已经完成 BF16 latent MLA，以及 `(TP, EP)` 为 `(1,1)`、`(2,1)`、 +`(4,1)`、`(1,2)`、`(1,4)`、`(2,2)`、`(4,2)`、`(4,4)` 时，vanilla、 +StreamingLLM、SnapKV、H2O、OmniKV、R-KV 与 decode CUDA Graph、Prefix +Cache 的完整交叉组合，硬件限定为 NVIDIA H100 80GB HBM3、`DP=1`。 +支持声明只覆盖这些实测组合，不向其他并行规模、其他 sparse 方法、prefix +offload 或多稀疏算法叠加外推。 ### 25.1 已落地的架构边界 @@ -1019,22 +1043,37 @@ HBM3。支持声明只覆盖这些实测组合。 `triton_h100`，执行失败不会静默切换实现。 - GLM MoE 复用 Qwen3-MoE 的 packed expert 物理执行，只保留 GLM 自己的 biased-sigmoid router、Dense/MoE topology 和 checkpoint 语义。 +- 联合 TP/EP 复用 Qwen3-MoE 的 outer-TP group topology：attention TP 为 + `T`、MoE EP 为 `E`、MoE TP 为 `T/E`、world size 为 `T`；纯 EP 保留 + replicated-attention world size `E`，纯 TP 保留普通 tensor parallel。 - Chat Completions 与 Responses API 均继续调用 Transformers response parser； 本地只提供 GLM 声明式 template。Terminal EOS 和 stop boundary 在 dispatcher / detokenizer 通用边界处理，不在 GLM parser 中写模型特判。 +- 稀疏方法继续复用 cache-manager-first 路径；MLA storage 提供 latent cache 的 + eviction、query-key materialization 和 graph-stable selection metadata，不在 + `attention.py` 增加 GLM 方法分支。 +- Decode CUDA Graph 保留 warmup graph 及其 private-pool ownership；业务请求用 + capture/replay counter 增量验证实际 replay，不通过销毁 warmup graph 后重捕获来 + 重置测量。 ### 25.2 验证结果 | 门禁 | 结果 | 证据 | | --- | --- | --- | -| CPU regression | `402 passed, 33 skipped, 87 subtests passed` | 完整 GLM/MLA/MoE/OpenAI/scheduler test selection | -| CUDA kernel/operator/model regression | `57 passed` | H100 kernel、operator、attention layer 与模型测试 | +| CPU regression | `1523 passed, 172 skipped, 239 subtests passed` | CPU-only 完整回归 | +| CUDA kernel/operator/model targeted regression | `83 passed` | H100 kernel、operator、attention layer、模型与稀疏 score lifecycle targeted selection | | Tiny multi-chunk TP1 | 17-token prompt、chunk size 8，实际 3 chunks；prefill 与两步 decode argmax 全匹配，max abs diff 分别为 `0.00390625`、`0.00390625`、`0.0048828125` | run manifest、raw comparison 与 stdout | | 真实 checkpoint TP1 | 48/48 shards、9491 tensors、55.77 GiB、8832 local expert shards、精确跳过 212 个 MTP tensors；prefill + 32 decode steps argmax 全匹配 | run manifest、raw comparison 与 stdout | | 真实 checkpoint TP2 | 每 rank 28.01 GiB，两个 rank 均加载 48/48 shards、9491 tensors；prefill + 32 decode steps 共 33/33 argmax 全匹配，greedy token 序列与 TP1 相同 | run manifest、raw comparison 与 stdout | | 真实 checkpoint TP4/EP1 | 每 rank 14.13 GiB，四个 rank 均绑定 `triton_h100` MLA 与 Triton MoE；短序列 33/33、257-token 五段 prefill + decode 9/9 argmax 全匹配 | run manifest、raw comparison 与 stdout | | TP4 synthetic API ragged + natural LCC | 64/256/768/1536-token synthetic API batch 连续三轮共 12/12 成功且跨轮逐字稳定；另以 1004/1023/1024/1025-token natural LCC 样本对 HF teacher-forced prefill 7/7、decode 27/28，唯一 teacher-forced argmax mismatch 为两个 HF 候选 logit 完全相等的 BF16 tie | raw/parsed responses、per-row comparison 与 request logs | | TP4 OpenAI client gates | Chat 普通/思考、Chat SSE、Responses、Responses SSE、function tool call、streaming tool call 和 4-way concurrency 共 8/8 client gates 通过；parser 实际委托 Transformers。服务在验证后由 operator 主动停止，teardown 不作为 client gate | command/env manifest、raw parser events、parsed responses、validation status 与 request logs | +| 真实 checkpoint EP2/EP4 | `TP=1, DP=1` 下 EP2 与 EP4 均成功；各 rank 的 router top-k、MoE 输出和最终 logits 精确一致，prefill 覆盖全部 expert shard | summary、完整 logits 与跨 rank comparison | +| 真实 checkpoint 联合 TP/EP topology | TP2/EP2、TP4/EP2、TP4/EP4 均成功；attention、expert、MoE-TP、data group 与 outer-TP 契约逐 rank 一致，所有 expert partition 在 prefill 实际命中，router top-k 与 MoE 输出跨 rank 精确一致 | per-rank group summary、expert hit count、router/MoE comparison 与 status | +| TP2/EP1 与 TP2/EP2 数值门禁 | 同 prompt/config 的 prefill logits max/mean/p99 abs diff 为 `1.5/0.174635/0.5`，final decode 为 `2.0/0.417368/1.0`；首 token 差异来自 reference 中两个候选的零 margin BF16 tie，后续 token 完全一致 | 两侧 float32 全词表 logits、token 序列、SHA256 与独立 comparison status | +| 真实 checkpoint Prefix Cache | vanilla/OmniKV radix 与 StreamingLLM/SnapKV/H2O/R-KV chain 共 6 种模式成功；每种均完成 2/2 请求，第二请求命中，eligible hit rate 为 1.0，chain 模式无 invalidation/prefix mismatch | aggregate metrics、resolved config 与 request logs | +| 真实 checkpoint decode CUDA Graph | vanilla、StreamingLLM、SnapKV、H2O、OmniKV、R-KV 共 6 种方法均通过 10/10 gates；两段 prompt 产生的 6 行全词表 logits shape 为 `[6, 154880]`，eager/Graph max abs diff 均为 `0.0`；每项 lifetime capture 为 1、业务 replay 增量为 6，且无 eager/force-eager fallback，方法特定 trigger 均有证据 | 完整 logits、runtime counter before/after/delta、selection/cache trace 与 status | +| Graph + Prefix 完整组合矩阵 | 8 种 `(TP, EP)` 布局 × 6 种方法共 48/48 case 成功；每项完成 2/2 请求且 eligible hit rate 为 1.0，每个 world rank 均 lifetime capture=1、业务 replay 增量=4、eager/force-eager 增量=0 | per-case metrics、per-rank Graph before/after/delta、prefix/chain lifecycle 与 aggregate status | 完整命令、依赖与设备版本、dirty status、raw/parsed output 和本机 artifact 路径保存在私有实验记录中，不写入公开仓库文档；上表只保留稳定的支持边界和 @@ -1049,9 +1088,13 @@ TP1、TP2 和 TP4 的真实 checkpoint 32 步 greedy token 序列全部一致， ### 25.3 尚未完成的门禁 -- CUDA Graph、Prefix Cache/offload、量化、MTP 和所有 sparse method。 +- Prefix Cache offload。 +- PyramidKV、QuEST、SkipKV、DeltaKV 等未列入当前 GLM 矩阵的方法。 +- 量化 checkpoint 与 MTP/speculative decoding。 - 128K/202K 长上下文容量与吞吐验证。 -- `TP=5` 等其他可整除配置，以及 `EP > 1`。 +- `TP=5` 等其他可整除配置与 `DP>1`。 +- SnapKV + H2O 等多稀疏算法叠加；当前“组合支持”指一个稀疏方法与 + TP/EP、CUDA Graph、Prefix Cache 的能力交叉，不定义新的混合稀疏算法。 - H100 以外的 GPU 架构。 这些组合在完成独立验证前必须继续 fail fast 或维持未支持声明。 diff --git a/docs/zh/features/supported-models.md b/docs/zh/features/supported-models.md index 21819784..af43594a 100644 --- a/docs/zh/features/supported-models.md +++ b/docs/zh/features/supported-models.md @@ -13,27 +13,32 @@ | Qwen3MoE | `qwen3_moe` | BF16 / FP16 / 块级 FP8 | ✅（TP > 1 时模型 dtype 仅支持 BF16） | 仅支持 1 | ✅ | | Qwen3.5 / Qwen3.6 | `qwen3_5` | BF16 / 块级 FP8 | ✅ | 仅支持 1 | 仅支持 1 | | Qwen3.6 MoE | `qwen3_5_moe` | BF16 / 块级 FP8 | ✅ | 仅支持 1 | ✅ | -| GLM-4.7-Flash | `glm4_moe_lite` | BF16 | 1 / 2 / 4（仅 H100） | 仅支持 1 | 仅支持 1 | +| GLM-4.7-Flash | `glm4_moe_lite` | BF16 | 1 / 2 / 4（仅 H100）⁵ | 仅支持 1 | 1 / 2 / 4⁵ | | Llama 3 / 3.1 | `llama` | BF16 / FP16 | ✅ | 仅支持 1 | 仅支持 1 | | MiniMax M2.7 | `minimax_m2` | 块级 FP8，非量化权重使用 BF16 | ✅ | 仅支持 1 | ✅ | TP 规模限制为 1 到 8，并且 checkpoint 维度（包括 attention head 数和 vocabulary 大小）必须能被所选 TP 规模整除。Qwen3MoE 的 EP 规模必须整除 `num_experts`；MiniMax M2.7 的 EP 规模必须整除 `num_local_experts`。 -当外层 TP 规模 `T` 大于 1 时，Qwen3MoE 和 MiniMax M2.7 使用混合并行布局： -attention TP 为 `T`、MoE EP 为 `E`、MoE TP 为 `T / E`，distributed world -size 为 `T`。该布局要求 `DP=1` 且 `T % E == 0`；专家数量必须能被 `E` -整除，MoE intermediate dimension 必须能被 `T / E` 整除。Qwen3MoE 的外层 -TP 要求模型 dtype 为 BF16；FP16 Qwen3MoE checkpoint 仅支持 `TP=1`。当 -`TP=1` 时，原有 EP 布局的 world size 为 `E`。 +当外层 TP 规模 `T` 大于 1 时，Qwen3MoE 和 MiniMax M2.7 使用混合并行布局； +GLM-4.7-Flash 在 `T > 1` 且 `E > 1` 时使用相同布局：attention TP 为 `T`、 +MoE EP 为 `E`、MoE TP 为 `T / E`，distributed world size 为 `T`。该布局 +要求 `DP=1` 且 `T % E == 0`；专家数量必须能被 `E` 整除，MoE intermediate +dimension 必须能被 `T / E` 整除。Qwen3MoE 的外层 TP 要求模型 dtype 为 +BF16；FP16 Qwen3MoE checkpoint 仅支持 `TP=1`。当 `TP=1` 时，原有 EP +布局的 world size 为 `E`。 块级 FP8 要求使用 E4M3 权重、动态激活量化以及 `128 x 128` 的权重块大小。 Qwen3.5/Qwen3.6 Dense 配置在内部统一规范为 `model_type=qwen3_5`；Qwen3.6 MoE 使用 `model_type=qwen3_5_moe`。 -GLM-4.7-Flash 当前仅支持 BF16、eager、vanilla latent MLA、`TP=1/2/4`、 -`DP=1` 和 `EP=1`，硬件限定为 NVIDIA H100 80GB HBM3。loader 会有意跳过 -checkpoint 中的 MTP 层。其他 GPU 架构、TP 规模、expert parallel、CUDA Graph、 -Prefix Cache、量化和稀疏方法必须经过独立验证后才会进入支持范围。 +GLM-4.7-Flash 在 NVIDIA H100 80GB HBM3 上使用 BF16 latent MLA，要求 +`DP=1` 且 `enforce_eager=True`。已验证的 `(TP, EP)` 布局为 `(1,1)`、 +`(2,1)`、`(4,1)`、`(1,2)`、`(1,4)`、`(2,2)`、`(4,2)` 和 `(4,4)`。 +在全部八种布局中，vanilla、StreamingLLM、SnapKV、H2O、OmniKV 和 R-KV +均支持 decode CUDA Graph 与 Prefix Cache 的联合组合。Prefix Cache 对 +vanilla 和 OmniKV 使用 radix 模式，对 StreamingLLM、SnapKV、H2O 和 R-KV +使用 chain 模式。Prefix offload、量化和其他稀疏方法仍不支持。loader +会有意跳过 checkpoint 中的 MTP 层。 ## 稀疏方法支持 @@ -44,7 +49,7 @@ Prefix Cache、量化和稀疏方法必须经过独立验证后才会进入支 | Qwen3MoE | ✅ | ✅ | ✅ | 实验性⁴ | ✅ | ✅ | ✅ | ✅ | — | — | | Qwen3.5 / Qwen3.6 | ✅ | ✅ | ✅ | 实验性⁴ | ✅ | ✅ | ✅ | ✅ | — | 匹配的 checkpoint³ | | Qwen3.6 MoE | ✅ | ✅ | ✅ | 实验性⁴ | ✅ | ✅ | ✅ | ✅ | — | — | -| GLM-4.7-Flash | ✅ | — | — | — | — | — | — | — | — | — | +| GLM-4.7-Flash | ✅⁵ | ✅⁵ | ✅⁵ | 实验性⁴⁵ | — | ✅⁵ | — | ✅⁵ | — | — | | Llama 3 / 3.1 | ✅ | ✅ | ✅ | 实验性⁴ | ✅ | ✅ | ✅ | ✅ | 指定 checkpoint¹ | 需要 compressor² | | MiniMax M2.7 | ✅ | ✅ | ✅ | 实验性⁴ | ✅ | ✅ | ✅ | ✅ | — | — | @@ -61,4 +66,8 @@ attention head 或 KV head 独立计算分数并保留 token，不跨 rank 聚 index。因此其算法行为不保证与 TP=1 或全局 head 选择等价；各模型原有的 TP、EP、DP 限制仍然适用。 +⁵ GLM 支持限定为上文列出的八种 `(TP, EP)` 布局，且要求 `DP=1`。在 +`TP>1` 时，基于 head 评分的稀疏方法使用 TP-local selection，不跨 rank +聚合 sparse index，因此其选择语义不保证与 `TP=1` 相同。 + `—` 表示当前不支持该组合。 diff --git a/scripts/benchmarks/bench_prefix_cache.py b/scripts/benchmarks/bench_prefix_cache.py index 9fc28f99..cd1c699d 100644 --- a/scripts/benchmarks/bench_prefix_cache.py +++ b/scripts/benchmarks/bench_prefix_cache.py @@ -42,6 +42,11 @@ "enable_prefix_caching": True, "label": "QuEST, prefix cache on", }, + "chain_streamingllm": { + "method": "streamingllm", + "enable_prefix_caching": True, + "label": "StreamingLLM, linear chain prefix cache", + }, "chain_snapkv": { "method": "snapkv", "enable_prefix_caching": True, @@ -77,6 +82,9 @@ "prefix_vanilla": "prefix_full", "omnikv": "prefix_omnikv", "quest": "prefix_quest", + "streamingllm": "chain_streamingllm", + "attention-sink": "chain_streamingllm", + "attention_sink": "chain_streamingllm", "snapkv": "chain_snapkv", "h2o": "chain_h2o", "pyramidkv": "chain_pyramidkv", @@ -461,6 +469,12 @@ def _case_engine_kwargs(args: argparse.Namespace, case_name: str, max_prompt_len "enforce_eager": True, "gpu_memory_utilization": float(args.gpu_memory_utilization), "tensor_parallel_size": int(args.tensor_parallel_size), + "expert_parallel_size": int( + getattr(args, "expert_parallel_size", 1) + ), + "decode_cuda_graph": bool( + getattr(args, "decode_cuda_graph", False) + ), "max_num_seqs_in_batch": int(args.max_active_requests), "max_decoding_seqs": int(args.max_active_requests), "max_num_batched_tokens": int(args.max_num_batched_tokens), @@ -495,6 +509,55 @@ def _case_engine_kwargs(args: argparse.Namespace, case_name: str, max_prompt_len return {key: value for key, value in hyper_params.items() if value is not None} +def _decode_graph_rank_summaries(llm: Any) -> list[dict[str, Any]]: + summaries = llm.debug_sparse_state_summaries() + graph_summaries: list[dict[str, Any]] = [] + for summary in summaries: + graph = summary.get("decode_cuda_graph") + if not isinstance(graph, dict): + raise RuntimeError( + "Worker debug summary is missing decode_cuda_graph evidence: " + f"world_rank={summary.get('world_rank')}." + ) + graph_summaries.append( + { + "world_rank": int(summary["world_rank"]), + **graph, + } + ) + return graph_summaries + + +def _decode_graph_rank_deltas( + before: list[dict[str, Any]], + after: list[dict[str, Any]], +) -> list[dict[str, Any]]: + before_by_rank = {int(item["world_rank"]): item for item in before} + after_by_rank = {int(item["world_rank"]): item for item in after} + if set(before_by_rank) != set(after_by_rank): + raise RuntimeError( + "Decode CUDA Graph evidence changed world ranks: " + f"before={sorted(before_by_rank)} after={sorted(after_by_rank)}." + ) + counter_names = ( + "capture_count", + "replay_count", + "eager_static_count", + "force_eager_count", + ) + return [ + { + "world_rank": rank, + **{ + name: int(after_by_rank[rank][name]) + - int(before_by_rank[rank][name]) + for name in counter_names + }, + } + for rank in sorted(before_by_rank) + ] + + def _token_vocab(tokenizer: Any) -> list[int]: special_ids = set(getattr(tokenizer, "all_special_ids", []) or []) vocab_values = sorted(set(int(token_id) for token_id in tokenizer.get_vocab().values())) @@ -1060,6 +1123,8 @@ def _summarize_records( cache_stats_after: dict[str, int], peak_memory_gb: float, elapsed_s: float, + decode_graph_before: list[dict[str, Any]] | None = None, + decode_graph_after: list[dict[str, Any]] | None = None, ) -> dict[str, Any]: success = [record for record in records if record.get("status") == "success"] failures = [record for record in records if record.get("status") != "success"] @@ -1092,6 +1157,43 @@ def _summarize_records( key: int(cache_stats_after.get(key, 0)) - int(cache_stats_before.get(key, 0)) for key in sorted(set(cache_stats_before) | set(cache_stats_after)) } + graph_required = bool(getattr(args, "decode_cuda_graph", False)) + graph_deltas = _decode_graph_rank_deltas( + decode_graph_before or [], + decode_graph_after or [], + ) + graph_failures: list[str] = [] + if graph_required: + if not graph_deltas: + graph_failures.append("missing per-rank CUDA Graph counters") + after_by_rank = { + int(item["world_rank"]): item + for item in (decode_graph_after or []) + } + for delta in graph_deltas: + rank = int(delta["world_rank"]) + after = after_by_rank[rank] + if int(after["capture_count"]) <= 0: + graph_failures.append( + f"world rank {rank} did not capture a decode graph" + ) + if int(delta["replay_count"]) <= 0: + graph_failures.append( + f"world rank {rank} did not replay a decode graph" + ) + if int(delta["eager_static_count"]) != 0: + graph_failures.append( + f"world rank {rank} used eager static decode" + ) + if int(delta["force_eager_count"]) != 0: + graph_failures.append( + f"world rank {rank} used forced eager fallback" + ) + if graph_failures: + summary_status = "metric_failed" + failure_status_counts["metric_failed"] = ( + failure_status_counts.get("metric_failed", 0) + 1 + ) by_turn: dict[str, dict[str, Any]] = {} for record in bench_success: @@ -1162,6 +1264,11 @@ def _summarize_records( "prefix_cache_stats_before": cache_stats_before, "prefix_cache_stats_after": cache_stats_after, "prefix_cache_stats_delta": stats_delta, + "decode_cuda_graph_required": graph_required, + "decode_cuda_graph_failures": graph_failures, + "decode_cuda_graph_before": decode_graph_before or [], + "decode_cuda_graph_after": decode_graph_after or [], + "decode_cuda_graph_delta": graph_deltas, "per_turn": per_turn, } return summary @@ -1200,6 +1307,7 @@ def _run_case_worker(case_name: str, args_dict: dict[str, Any], case_dir: str) - started_s = time.perf_counter() llm = LLM(args.model_path, **engine_kwargs) cache_stats_before = _cache_stats(llm) + decode_graph_before = _decode_graph_rank_summaries(llm) records: list[dict[str, Any]] = [] workloads = set(_split_csv(args.workloads)) @@ -1237,6 +1345,7 @@ def _run_case_worker(case_name: str, args_dict: dict[str, Any], case_dir: str) - torch.cuda.max_memory_allocated() / (1024**3) if torch.cuda.is_available() else 0.0 ) cache_stats_after = _cache_stats(llm) + decode_graph_after = _decode_graph_rank_summaries(llm) summary = _summarize_records( case_name=case_name, case_config=CASE_PRESETS[case_name], @@ -1247,6 +1356,8 @@ def _run_case_worker(case_name: str, args_dict: dict[str, Any], case_dir: str) - cache_stats_after=cache_stats_after, peak_memory_gb=peak_memory_gb, elapsed_s=elapsed_s, + decode_graph_before=decode_graph_before, + decode_graph_after=decode_graph_after, ) (case_dir_path / "aggregate_metrics.json").write_text( json.dumps(summary, indent=2, ensure_ascii=False) + "\n", @@ -1414,6 +1525,12 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--gpu_memory_utilization", type=float, default=0.65) parser.add_argument("--tensor_parallel_size", type=int, default=1) + parser.add_argument("--expert_parallel_size", type=int, default=1) + parser.add_argument( + "--decode_cuda_graph", + action=argparse.BooleanOptionalAction, + default=False, + ) parser.add_argument("--max_active_requests", type=int, default=4) parser.add_argument("--max_num_batched_tokens", type=int, default=8192) parser.add_argument("--chunk_prefill_size", type=int, default=4096) diff --git a/scripts/debug/compare_decode_graph_eager_logits.py b/scripts/debug/compare_decode_graph_eager_logits.py index 23e0bb92..73ed15cb 100644 --- a/scripts/debug/compare_decode_graph_eager_logits.py +++ b/scripts/debug/compare_decode_graph_eager_logits.py @@ -2,6 +2,7 @@ import argparse import gc +import hashlib import json import multiprocessing as mp import os @@ -13,6 +14,41 @@ import torch +METHOD_CHOICES = ( + "vanilla", + "streamingllm", + "attention-sink", + "attention_sink", + "snapkv", + "pyramidkv", + "h2o", + "rkv", + "skipkv", + "quest", + "omnikv", + "deltakv", + "deltakv-less-memory", + "deltakv-less-memory-cudagraph", +) + +GLM_GRAPH_METHODS = frozenset( + {"vanilla", "streamingllm", "snapkv", "h2o", "omnikv", "rkv"} +) + + +def _tensor_sha256(tensor: torch.Tensor) -> str: + raw = tensor.detach().contiguous().cpu().view(torch.uint8).numpy().tobytes() + return hashlib.sha256(raw).hexdigest() + + +def _file_sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + def _load_json_arg(value: str) -> dict[str, Any]: if value is None: return {} @@ -37,12 +73,27 @@ def _topk_overlap(a: torch.Tensor, b: torch.Tensor, k: int) -> dict[str, float | return {"intersection": intersection, "ratio": float(intersection / k if k else 1.0)} -def _compare_logits(eager: torch.Tensor, graph: torch.Tensor) -> dict[str, Any]: +def _compare_logits( + eager: torch.Tensor, + graph: torch.Tensor, + *, + atol: float = 0.05, + rtol: float = 0.05, +) -> dict[str, Any]: if eager.shape != graph.shape: - raise ValueError(f"Logit shape mismatch: eager={tuple(eager.shape)} graph={tuple(graph.shape)}") + raise ValueError( + "Logit shape mismatch: " + f"eager={tuple(eager.shape)} graph={tuple(graph.shape)}" + ) diff = (eager - graph).abs() + tolerance = float(atol) + float(rtol) * eager.abs() + tolerance_ratio = diff / tolerance.clamp_min(torch.finfo(torch.float32).eps) result: dict[str, Any] = { "shape": list(eager.shape), + "atol": float(atol), + "rtol": float(rtol), + "within_tolerance": bool(torch.all(diff <= tolerance).item()), + "max_tolerance_ratio": float(tolerance_ratio.max().item()), "max_abs_diff": float(diff.max().item()), "mean_abs_diff": float(diff.mean().item()), "argmax_match": eager.argmax(dim=-1).tolist() == graph.argmax(dim=-1).tolist(), @@ -87,11 +138,19 @@ def _tensor_summary(tensor: torch.Tensor | None, *, limit: int = 16) -> dict[str "shape": [int(x) for x in detached.shape], "dtype": str(detached.dtype), "numel": int(flat.numel()), + "sha256": _tensor_sha256(detached), } if flat.numel() == 0: out.update({"sum": 0, "min": None, "max": None, "preview": []}) return out - if detached.dtype in (torch.int8, torch.int16, torch.int32, torch.int64, torch.long, torch.bool): + if detached.dtype in ( + torch.int8, + torch.int16, + torch.int32, + torch.int64, + torch.long, + torch.bool, + ): values = flat.to(torch.int64) out.update( { @@ -114,7 +173,189 @@ def _tensor_summary(tensor: torch.Tensor | None, *, limit: int = 16) -> dict[str return out -def _capture_selection_trace(llm, *, step: int, use_graph: bool) -> dict[str, Any]: +def _canonical_method(method: str) -> str: + aliases = { + "attention-sink": "streamingllm", + "attention_sink": "streamingllm", + } + return aliases.get(str(method), str(method)) + + +def _install_method_instrumentation(llm) -> dict[str, int]: + """Count method-path calls made after engine warmup. + + These wrappers are diagnostic-only and live inside the isolated worker. + They make post-forward eviction/compaction calls auditable without changing + the cache-manager hot path or relying on log text. + """ + + calls: dict[str, int] = {} + targets = ( + ( + "cache", + llm.model_runner.cache_manager, + ( + "free_prefix_recent_slots_batch_layers", + "free_part_slots_batch_layers", + "free_part_slots_batch", + "free_part_slots", + "evict_after_decode", + "update_decode_attention_scores_all_layers", + "rkv_query_attention_scores_batch", + "rkv_query_attention_scores", + "select_rkv_indices_batch", + "select_rkv_indices", + "materialize_attention_keys", + ), + ), + ( + "controller", + llm.model_runner.sparse_controller, + ("_update_dynamic_omnikv_indices",), + ), + ) + for prefix, target, names in targets: + for name in names: + original = getattr(target, name, None) + if not callable(original): + continue + key = f"{prefix}.{name}" + calls[key] = 0 + + def wrapped(*args, _key=key, _original=original, **kwargs): + calls[_key] += 1 + return _original(*args, **kwargs) + + setattr(target, name, wrapped) + return calls + + +def _graph_counter_snapshot(runner) -> dict[str, int]: + return { + "capture_count": int(runner.capture_count), + "replay_count": int(runner.replay_count), + "eager_static_count": int(runner.eager_static_count), + "force_eager_count": int(runner.force_eager_count), + "graph_count": sum( + state.graph is not None for state in runner._graphs.values() + ), + } + + +def _start_graph_measurement(llm) -> dict[str, int]: + """Snapshot warmup counters without invalidating its CUDA graph pool. + + A captured graph owns the allocator private pool represented by + ``runner.graph_pool``. Clearing the last graph while a captured allocation + is still referenced leaves that handle non-reusable in PyTorch. Preserve + the warmup graphs and report business-request activity as counter deltas. + """ + + runner = llm.model_runner.decode_cuda_graph_runner + return _graph_counter_snapshot(runner) + + +def _graph_runtime_summary( + llm, + *, + use_graph: bool, + counters_before: dict[str, int], +) -> dict[str, Any]: + runner = llm.model_runner.decode_cuda_graph_runner + graph_states = [state for state in runner._graphs.values() if state.graph is not None] + counters_after = _graph_counter_snapshot(runner) + counter_delta = { + name: int(counters_after[name]) - int(counters_before[name]) + for name in ( + "capture_count", + "replay_count", + "eager_static_count", + "force_eager_count", + ) + } + return { + "requested": bool(use_graph), + "config_enabled": bool(llm.config.decode_cuda_graph), + "model": str(llm.config.model), + "model_type": str(getattr(llm.config.hf_config, "model_type", "")), + "configured_sparse_method": str( + getattr(llm.config, "vllm_sparse_method", "") or "vanilla" + ), + "runner_method": str(runner.method or "vanilla"), + "graph_active": bool(graph_states), + "graph_count": len(graph_states), + "capture_count": int(runner.capture_count), + "replay_count": int(runner.replay_count), + "eager_static_count": int(runner.eager_static_count), + "force_eager_count": int(runner.force_eager_count), + "counters_before": dict(counters_before), + "counters_after": counters_after, + "counter_delta": counter_delta, + "fallback": bool( + counter_delta["force_eager_count"] + or (use_graph and not graph_states) + ), + "graph_keys": [ + { + "method": str(state.key.method or "vanilla"), + "batch_size": int(state.key.batch_size), + "context_capacity": int(state.key.context_capacity), + "is_long_text": bool(state.key.is_long_text), + "capture_sampling": bool(state.key.capture_sampling), + } + for state in graph_states + ], + } + + +def _validate_graph_runtime(summary: dict[str, Any]) -> None: + failures = [] + delta = summary.get("counter_delta", summary) + if not summary.get("config_enabled"): + failures.append("decode_cuda_graph config is disabled") + if not summary.get("graph_active") or int(summary.get("graph_count", 0)) <= 0: + failures.append("no captured CUDA Graph is active") + if int(summary.get("capture_count", 0)) <= 0: + failures.append("no CUDA Graph capture exists in the engine lifetime") + if int(delta.get("replay_count", 0)) <= 0: + failures.append("business-request replay_count did not increase") + if int(delta.get("eager_static_count", 0)) != 0: + failures.append("graph run executed eager-static decode") + if int(delta.get("force_eager_count", 0)) != 0: + failures.append("graph run forced eager decode") + if summary.get("fallback"): + failures.append("graph runtime reported fallback") + if failures: + raise RuntimeError("CUDA Graph runtime gate failed: " + "; ".join(failures)) + + +def _validate_eager_runtime(summary: dict[str, Any]) -> None: + failures = [] + delta = summary.get("counter_delta", summary) + if summary.get("config_enabled"): + failures.append("eager control unexpectedly enabled decode_cuda_graph") + if summary.get("graph_active") or int(summary.get("graph_count", 0)) != 0: + failures.append("eager control retained a captured CUDA Graph") + if int(delta.get("eager_static_count", 0)) <= 0: + failures.append("eager control did not execute eager-static decode") + if int(delta.get("capture_count", 0)) != 0 or int( + delta.get("replay_count", 0) + ) != 0: + failures.append("eager control captured or replayed a CUDA Graph") + if int(delta.get("force_eager_count", 0)) != 0: + failures.append("eager control unexpectedly used force-eager routing") + if failures: + raise RuntimeError("Eager runtime gate failed: " + "; ".join(failures)) + + +def _capture_selection_trace( + llm, + *, + step: int, + stage: str, + logical_context_len: int, + use_graph: bool, +) -> dict[str, Any]: sparse_controller = llm.model_runner.sparse_controller cache_manager = llm.model_runner.cache_manager layers: dict[str, Any] = {} @@ -123,28 +364,203 @@ def _capture_selection_trace(llm, *, step: int, use_graph: bool) -> dict[str, An attn_score = state.attn_score should_record = ( active is not None + or state.active_indices is not None + or state.active_slots is not None or attn_score is not None or int(layer_idx) in set(getattr(sparse_controller, "obs_layer_ids", [])) ) if not should_record: continue layers[str(int(layer_idx))] = { + "active_indices": _tensor_summary(state.active_indices), + "active_slots": _tensor_summary(state.active_slots), "active_compressed_indices": _tensor_summary(active), "attn_score": _tensor_summary(attn_score, limit=8), "context_lens": _tensor_summary(state.context_lens), "req_indices": _tensor_summary(state.req_indices), - "max_context_len": None if state.max_context_len is None else int(state.max_context_len), + "max_context_len": ( + None + if state.max_context_len is None + else int(state.max_context_len) + ), } compressed_lens = getattr(cache_manager, "_deltakv_decode_static_compressed_lens", None) + rkv_materializer_layers: list[int] = [] + layer_indices = getattr(cache_manager, "kv_transformer_layer_indices", lambda: ())() + has_materializer = getattr(cache_manager, "has_attention_key_materializer", None) + if callable(has_materializer): + rkv_materializer_layers = [ + int(layer_idx) + for layer_idx in layer_indices + if has_materializer(int(layer_idx)) + ] return { "step": int(step), + "stage": str(stage), + "logical_context_len": int(logical_context_len), "use_graph": bool(use_graph), "compressed_lens": _tensor_summary(compressed_lens), "layers": layers, + "dynamic_selection": sparse_controller.debug_state_summary()[ + "dynamic_selection" + ], + "cache": cache_manager.debug_state_summary(), + "rkv_materializer_layers": rkv_materializer_layers, } +def _live_row_lengths(trace: dict[str, Any]) -> list[int]: + live_rows = trace.get("cache", {}).get("live_rows", {}) + return [ + int(record["row_len"]) + for records in live_rows.values() + for record in records + ] + + +def _has_physical_compaction(traces: list[dict[str, Any]]) -> bool: + for trace in traces: + row_lengths = _live_row_lengths(trace) + if row_lengths and min(row_lengths) < int(trace["logical_context_len"]): + return True + return False + + +def _has_omnikv_selection(traces: list[dict[str, Any]]) -> bool: + for trace in traces: + logical_context_len = int(trace["logical_context_len"]) + for layer in trace.get("layers", {}).values(): + active_slots = layer.get("active_slots") + context_lens = layer.get("context_lens") + if ( + active_slots is not None + and int(active_slots.get("numel", 0)) > 0 + and context_lens is not None + and context_lens.get("max") is not None + and int(context_lens["max"]) < logical_context_len + ): + return True + return False + + +def _latest_h2o_state(traces: list[dict[str, Any]]) -> dict[str, Any]: + for trace in reversed(traces): + h2o = trace.get("cache", {}).get("h2o") + if isinstance(h2o, dict): + return h2o + return {} + + +def _build_method_trigger_evidence( + method: str, + traces: list[dict[str, Any]], + method_calls: dict[str, int], +) -> dict[str, Any]: + method = _canonical_method(method) + physical_compaction = _has_physical_compaction(traces) + positive_calls = {key: int(value) for key, value in method_calls.items() if int(value) > 0} + evidence: dict[str, Any] = { + "method": method, + "required": method in GLM_GRAPH_METHODS - {"vanilla"}, + "triggered": method == "vanilla", + "trigger_kind": "dense_baseline" if method == "vanilla" else "", + "physical_compaction": physical_compaction, + "method_calls": positive_calls, + } + if method in {"streamingllm", "snapkv"}: + compaction_calls = sum( + count + for name, count in positive_calls.items() + if "free_" in name + ) + evidence.update( + triggered=bool(physical_compaction and compaction_calls > 0), + trigger_kind="physical_eviction_compaction", + compaction_call_count=int(compaction_calls), + ) + elif method == "h2o": + h2o = _latest_h2o_state(traces) + counters = h2o.get("counters", {}) + eviction_count = sum( + int(counters.get(name, 0)) + for name in ( + "intermediate_prefill_evictions", + "final_prefill_evictions", + "decode_evictions", + ) + ) + dropped_tokens = int(counters.get("dropped_tokens", 0)) + ring_fallback_rows = int( + h2o.get("ring_counters", {}).get("fallback_rows", 0) + ) + evidence.update( + triggered=bool(eviction_count > 0 and dropped_tokens > 0), + trigger_kind="score_eviction_compaction", + h2o_counters=counters, + h2o_ring_counters=h2o.get("ring_counters", {}), + eviction_count=eviction_count, + dropped_tokens=dropped_tokens, + internal_fallback_rows=ring_fallback_rows, + ) + elif method == "omnikv": + selection_calls = int( + positive_calls.get("controller._update_dynamic_omnikv_indices", 0) + ) + selected = _has_omnikv_selection(traces) + captured_replay = any(bool(trace.get("use_graph")) for trace in traces) + evidence.update( + triggered=bool( + selected and (selection_calls > 0 or captured_replay) + ), + trigger_kind="dynamic_topk_selection", + selection_call_count=selection_calls, + compressed_selection_observed=selected, + execution_mode=( + "captured_replay" + if captured_replay and selection_calls == 0 + else "python_eager_or_capture" + ), + ) + elif method == "rkv": + score_calls = sum( + count + for name, count in positive_calls.items() + if "rkv_query_attention_scores" in name + ) + materializer_calls = int( + positive_calls.get("cache.materialize_attention_keys", 0) + ) + materializer_layers = sorted( + { + int(layer_idx) + for trace in traces + for layer_idx in trace.get("rkv_materializer_layers", []) + } + ) + evidence.update( + triggered=bool( + physical_compaction + and score_calls > 0 + and materializer_calls > 0 + and materializer_layers + ), + trigger_kind="query_scored_eviction_compaction", + query_score_call_count=int(score_calls), + materializer_call_count=materializer_calls, + materializer_layers=materializer_layers, + ) + return evidence + + +def _validate_method_trigger(evidence: dict[str, Any]) -> None: + if evidence.get("required") and not evidence.get("triggered"): + raise RuntimeError( + "Sparse method trigger gate failed: " + + json.dumps(evidence, sort_keys=True) + ) + + def _run_decode_logits( *, model_path: str, @@ -155,7 +571,7 @@ def _run_decode_logits( hyper_params: dict[str, Any], use_graph: bool, trace_selection: bool = False, -) -> tuple[torch.Tensor, list[dict[str, Any]]]: +) -> tuple[torch.Tensor, list[dict[str, Any]], dict[str, Any]]: from sparsevllm import LLM, SamplingParams if os.getenv("SPARSEVLLM_DEBUG_SKIP_ENGINE_WARMUP", "0") == "1": @@ -172,9 +588,12 @@ def _run_decode_logits( "throughput_log_interval_s": 0.0, } llm = LLM(model_path, **engine_kwargs) + graph_counters_before = _start_graph_measurement(llm) + method_calls = _install_method_instrumentation(llm) captured: list[torch.Tensor] = [] trace: list[dict[str, Any]] = [] - decode_step = 0 + runtime_step = 0 + generated_token_outputs: list[dict[str, Any]] = [] if not use_graph: runner = llm.model_runner @@ -189,9 +608,22 @@ def wrapped_run_model(input_ids, positions, is_prefill): runner.run_model = wrapped_run_model if runner.decode_cuda_graph_runner is not None: runner.decode_cuda_graph_runner.run_model = wrapped_run_model + else: + graph_runner = llm.model_runner.decode_cuda_graph_runner + original_graph_run = graph_runner.run + + def wrapped_graph_run(*args, **kwargs): + logits, token_ids = original_graph_run(*args, **kwargs) + if logits is not None: + captured.append(logits.detach().float().cpu()) + return logits, token_ids + + graph_runner.run = wrapped_graph_run try: for round_idx, prompt_len in enumerate(prompt_lens): + round_decode_step = 0 + round_prefilled_tokens = 0 prompt_token_ids = [] for batch_idx in range(batch_size): # Use deterministic non-uniform prompts so sparse selection and @@ -207,20 +639,63 @@ def wrapped_run_model(input_ids, positions, is_prefill): while not llm.is_finished(): _, num_tokens = llm.step() - if num_tokens < 0: - decode_step += 1 - if use_graph: - runner = llm.model_runner.decode_cuda_graph_runner - if runner is None: - raise RuntimeError("decode_cuda_graph runner was not initialized.") - if runner.last_state_key is None or runner.last_real_batch_size is None: - raise RuntimeError("No graph logits were captured.") - state = runner._graphs[runner.last_state_key] - if state.logits is None: - raise RuntimeError("Last graph state has no logits.") - captured.append(state.logits[:runner.last_real_batch_size].detach().float().cpu()) - if trace_selection: - trace.append(_capture_selection_trace(llm, step=decode_step, use_graph=use_graph)) + runtime_step += 1 + if num_tokens > 0: + round_prefilled_tokens += int(num_tokens) // int(batch_size) + stage = "prefill" + logical_context_len = min( + int(prompt_len), + int(round_prefilled_tokens), + ) + elif num_tokens < 0: + round_decode_step += 1 + stage = "decode" + logical_context_len = int(prompt_len) + int(round_decode_step) + else: + stage = "idle" + logical_context_len = int(prompt_len) + int(round_decode_step) + + generated_token_outputs.append( + { + "step": int(runtime_step), + "round": int(round_idx), + "stage": stage, + "token_outputs": [ + { + "seq_id": int(seq_id), + "token_ids": [int(token_id) for token_id in token_ids], + } + for seq_id, token_ids in llm.last_step_token_outputs + ], + } + ) + if num_tokens != 0: + trace.append( + _capture_selection_trace( + llm, + step=runtime_step, + stage=stage, + logical_context_len=logical_context_len, + use_graph=use_graph, + ) + ) + + graph_runtime = _graph_runtime_summary( + llm, + use_graph=use_graph, + counters_before=graph_counters_before, + ) + method_evidence = _build_method_trigger_evidence( + method, + trace, + method_calls, + ) + runtime_evidence = { + "graph": graph_runtime, + "method_trigger": method_evidence, + "method_calls": {key: int(value) for key, value in method_calls.items()}, + "generated_token_outputs": generated_token_outputs, + } finally: llm.exit() del llm @@ -229,19 +704,31 @@ def wrapped_run_model(input_ids, positions, is_prefill): if not captured: raise RuntimeError("No decode logits captured. Use max_tokens >= 3.") - return torch.cat(captured, dim=0), trace + del trace_selection + return torch.cat(captured, dim=0), trace, runtime_evidence def _run_decode_logits_worker(result_queue, kwargs: dict[str, Any]): try: - logits, trace = _run_decode_logits(**kwargs) - result_queue.put(("ok", {"logits": logits.numpy(), "trace": trace})) + logits, trace, runtime = _run_decode_logits(**kwargs) + result_queue.put( + ( + "ok", + { + "logits": logits.numpy(), + "trace": trace, + "runtime": runtime, + }, + ) + ) except BaseException: result_queue.put(("error", traceback.format_exc())) raise -def _run_decode_logits_isolated(**kwargs) -> tuple[torch.Tensor, list[dict[str, Any]]]: +def _run_decode_logits_isolated( + **kwargs, +) -> tuple[torch.Tensor, list[dict[str, Any]], dict[str, Any]]: ctx = mp.get_context("spawn") result_queue = ctx.Queue(maxsize=1) process = ctx.Process(target=_run_decode_logits_worker, args=(result_queue, kwargs)) @@ -254,31 +741,29 @@ def _run_decode_logits_isolated(**kwargs) -> tuple[torch.Tensor, list[dict[str, raise TimeoutError("Timed out waiting for decode logits worker.") from exc process.join() if process.exitcode != 0 or status != "ok": - raise RuntimeError(f"Decode logits worker failed with exitcode={process.exitcode}:\n{payload}") - return torch.from_numpy(payload["logits"]), payload["trace"] + raise RuntimeError( + "Decode logits worker failed with " + f"exitcode={process.exitcode}:\n{payload}" + ) + return ( + torch.from_numpy(payload["logits"]), + payload["trace"], + payload["runtime"], + ) -def main(): - parser = argparse.ArgumentParser(description="Compare Sparse-VLLM eager decode logits with decode CUDA Graph logits.") +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description=( + "Compare full eager and decode-CUDA-Graph logits and prove that " + "the requested sparse runtime path executed." + ) + ) parser.add_argument("--model_path", required=True) parser.add_argument( "--method", default="vanilla", - choices=( - "vanilla", - "streamingllm", - "attention-sink", - "attention_sink", - "snapkv", - "pyramidkv", - "rkv", - "skipkv", - "quest", - "omnikv", - "deltakv", - "deltakv-less-memory", - "deltakv-less-memory-cudagraph", - ), + choices=METHOD_CHOICES, ) parser.add_argument("--prompt_len", type=int, default=2048) parser.add_argument( @@ -291,55 +776,217 @@ def main(): parser.add_argument("--max_tokens", type=int, default=3) parser.add_argument("--hyper_params", default="{}") parser.add_argument("--output", required=True) + parser.add_argument("--atol", type=float, default=0.05) + parser.add_argument("--rtol", type=float, default=0.05) parser.add_argument("--trace_selection", action="store_true") - args = parser.parse_args() - - if args.max_tokens < 3: - raise ValueError("--max_tokens must be >= 3 to force at least one decode step.") - - hyper_params = _load_json_arg(args.hyper_params) - prompt_lens = [args.prompt_len] - if args.second_prompt_len is not None: - prompt_lens.append(args.second_prompt_len) - eager_logits, eager_trace = _run_decode_logits_isolated( - model_path=args.model_path, - method=args.method, - prompt_lens=prompt_lens, - batch_size=args.batch_size, - max_tokens=args.max_tokens, - hyper_params=hyper_params, - use_graph=False, - trace_selection=args.trace_selection, - ) - graph_logits, graph_trace = _run_decode_logits_isolated( - model_path=args.model_path, - method=args.method, - prompt_lens=prompt_lens, - batch_size=args.batch_size, - max_tokens=args.max_tokens, - hyper_params=hyper_params, - use_graph=True, - trace_selection=args.trace_selection, - ) + return parser + + +def _generated_token_ids(runtime: dict[str, Any]) -> list[int]: + return [ + int(token_id) + for step in runtime["generated_token_outputs"] + for record in step["token_outputs"] + for token_id in record["token_ids"] + ] + - output = { - "status": "success", - "method": args.method, - "prompt_lens": prompt_lens, - "batch_size": args.batch_size, - "max_tokens": args.max_tokens, - "hyper_params": hyper_params, - "comparison": _compare_logits(eager_logits, graph_logits), +def _save_full_logits_artifact( + path: Path, + *, + eager: torch.Tensor, + graph: torch.Tensor, +) -> dict[str, Any]: + path.parent.mkdir(parents=True, exist_ok=True) + torch.save( + { + "scope": "all_decode_rows_and_full_vocabulary", + "eager": eager.contiguous(), + "graph": graph.contiguous(), + }, + path, + ) + return { + "path": str(path.resolve()), + "scope": "all_decode_rows_and_full_vocabulary", + "artifact_sha256": _file_sha256(path), + "eager": _tensor_summary(eager), + "graph": _tensor_summary(graph), } - if args.trace_selection: - output["selection_trace"] = { - "eager": eager_trace, - "graph": graph_trace, - } + + +def _validation_error(validator, payload: dict[str, Any]) -> str | None: + try: + validator(payload) + except RuntimeError as exc: + return str(exc) + return None + + +def main(argv: list[str] | None = None): + args = _build_parser().parse_args(argv) output_path = Path(args.output) output_path.parent.mkdir(parents=True, exist_ok=True) - output_path.write_text(json.dumps(output, indent=2), encoding="utf-8") - print(json.dumps(output["comparison"], indent=2)) + wrote_current_output = False + + try: + if args.max_tokens < 3: + raise ValueError( + "--max_tokens must be >= 3 to force at least one decode step." + ) + if args.prompt_len <= 0 or args.batch_size <= 0: + raise ValueError("--prompt_len and --batch_size must be positive.") + if args.second_prompt_len is not None and args.second_prompt_len <= 0: + raise ValueError("--second_prompt_len must be positive when provided.") + if args.atol < 0 or args.rtol < 0: + raise ValueError("--atol and --rtol must be non-negative.") + + hyper_params = _load_json_arg(args.hyper_params) + prompt_lens = [args.prompt_len] + if args.second_prompt_len is not None: + prompt_lens.append(args.second_prompt_len) + eager_logits, eager_trace, eager_runtime = _run_decode_logits_isolated( + model_path=args.model_path, + method=args.method, + prompt_lens=prompt_lens, + batch_size=args.batch_size, + max_tokens=args.max_tokens, + hyper_params=hyper_params, + use_graph=False, + trace_selection=args.trace_selection, + ) + graph_logits, graph_trace, graph_runtime = _run_decode_logits_isolated( + model_path=args.model_path, + method=args.method, + prompt_lens=prompt_lens, + batch_size=args.batch_size, + max_tokens=args.max_tokens, + hyper_params=hyper_params, + use_graph=True, + trace_selection=args.trace_selection, + ) + + logits_artifact_path = output_path.with_name( + output_path.stem + ".full_logits.pt" + ) + logits_artifact = _save_full_logits_artifact( + logits_artifact_path, + eager=eager_logits, + graph=graph_logits, + ) + comparison = _compare_logits( + eager_logits, + graph_logits, + atol=args.atol, + rtol=args.rtol, + ) + eager_token_ids = _generated_token_ids(eager_runtime) + graph_token_ids = _generated_token_ids(graph_runtime) + token_ids_match = eager_token_ids == graph_token_ids + eager_runtime_error = _validation_error( + _validate_eager_runtime, + eager_runtime["graph"], + ) + graph_runtime_error = _validation_error( + _validate_graph_runtime, + graph_runtime["graph"], + ) + eager_method_error = _validation_error( + _validate_method_trigger, + eager_runtime["method_trigger"], + ) + graph_method_error = _validation_error( + _validate_method_trigger, + graph_runtime["method_trigger"], + ) + gates = { + "full_logits_within_tolerance": bool( + comparison["within_tolerance"] + ), + "argmax_match": bool(comparison["argmax_match"]), + "generated_token_ids_match": token_ids_match, + "eager_runtime_contract": eager_runtime_error is None, + "graph_runtime_contract": graph_runtime_error is None, + "graph_capture_observed": bool( + graph_runtime["graph"]["capture_count"] > 0 + ), + "graph_replay_observed": bool( + graph_runtime["graph"]["counter_delta"]["replay_count"] > 0 + ), + "no_eager_or_force_eager_fallback": bool( + graph_runtime["graph"]["counter_delta"]["eager_static_count"] + == 0 + and graph_runtime["graph"]["counter_delta"][ + "force_eager_count" + ] + == 0 + and not graph_runtime["graph"]["fallback"] + ), + "eager_method_triggered": bool( + eager_method_error is None + ), + "graph_method_triggered": bool( + graph_method_error is None + ), + } + passed = all(gates.values()) + output = { + "status": "success" if passed else "failed", + "method": args.method, + "prompt_lens": prompt_lens, + "batch_size": args.batch_size, + "max_tokens": args.max_tokens, + "hyper_params": hyper_params, + "comparison": comparison, + "generated_token_ids": { + "eager": eager_token_ids, + "graph": graph_token_ids, + "match": token_ids_match, + }, + "full_logits_artifact": logits_artifact, + "runtime": { + "eager": eager_runtime, + "graph": graph_runtime, + }, + "gates": gates, + "gate_errors": { + key: value + for key, value in { + "eager_runtime_contract": eager_runtime_error, + "graph_runtime_contract": graph_runtime_error, + "eager_method_triggered": eager_method_error, + "graph_method_triggered": graph_method_error, + }.items() + if value is not None + }, + } + if args.trace_selection: + output["selection_trace"] = { + "eager": eager_trace, + "graph": graph_trace, + } + output_path.write_text(json.dumps(output, indent=2), encoding="utf-8") + wrote_current_output = True + print(json.dumps({"status": output["status"], "gates": gates}, indent=2)) + if not passed: + raise RuntimeError( + "Eager-vs-CUDA-Graph validation gates failed: " + + json.dumps(gates, sort_keys=True) + ) + except BaseException: + if not wrote_current_output: + output_path.write_text( + json.dumps( + { + "status": "failed", + "method": args.method, + "error": traceback.format_exc(), + }, + indent=2, + ), + encoding="utf-8", + ) + raise if __name__ == "__main__": diff --git a/scripts/validation/compare_glm_parallel_variants.py b/scripts/validation/compare_glm_parallel_variants.py new file mode 100644 index 00000000..9d62eca2 --- /dev/null +++ b/scripts/validation/compare_glm_parallel_variants.py @@ -0,0 +1,293 @@ +#!/usr/bin/env python3 +"""Compare two GLM checkpoint parallel-validation artifacts.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import sys +import traceback +from pathlib import Path +from typing import Any + +import torch + + +LOGIT_KEYS = ("prefill", "final_decode") +MATCHED_RUN_FIELDS = ( + "model_path", + "seed", + "prompt_len", + "prompt_sha256", + "max_tokens", + "sparse_method", + "tensor_parallel_size", +) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Compare real-checkpoint GLM parallel validation artifacts." + ) + parser.add_argument("--reference_dir", required=True) + parser.add_argument("--candidate_dir", required=True) + parser.add_argument("--output", required=True) + parser.add_argument("--max_abs_diff", type=float, required=True) + parser.add_argument("--mean_abs_diff", type=float, required=True) + parser.add_argument("--p99_abs_diff", type=float, required=True) + parser.add_argument( + "--allow_first_token_tie", + action=argparse.BooleanOptionalAction, + default=False, + help=( + "Accept a first generated-token mismatch when the candidate token " + "is tied for the maximum in the reference prefill logits and all " + "later generated tokens match." + ), + ) + parser.add_argument( + "--tie_atol", + type=float, + default=0.0, + help="Absolute tolerance for the reference-logit tie check.", + ) + return parser.parse_args() + + +def _load_json(path: Path) -> dict[str, Any]: + payload = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(payload, dict): + raise TypeError(f"Expected a JSON object in {path}.") + return payload + + +def _file_sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _load_artifact(run_dir: Path) -> tuple[dict[str, Any], dict[str, torch.Tensor]]: + summary_path = run_dir / "summary.json" + logits_path = run_dir / "logits.pt" + if not summary_path.is_file() or not logits_path.is_file(): + raise FileNotFoundError( + f"Expected summary.json and logits.pt under {run_dir}." + ) + summary = _load_json(summary_path) + if summary.get("status") != "success": + raise RuntimeError( + f"Parallel validation did not succeed in {run_dir}: " + f"status={summary.get('status')!r}." + ) + raw_logits = torch.load(logits_path, map_location="cpu", weights_only=True) + if not isinstance(raw_logits, dict): + raise TypeError(f"Expected a tensor dictionary in {logits_path}.") + logits: dict[str, torch.Tensor] = {} + for key in LOGIT_KEYS: + tensor = raw_logits.get(key) + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"Missing tensor {key!r} in {logits_path}.") + logits[key] = tensor.detach().float().cpu().contiguous() + return summary, logits + + +def _compare_tensor( + reference: torch.Tensor, + candidate: torch.Tensor, + *, + max_abs_limit: float, + mean_abs_limit: float, + p99_abs_limit: float, +) -> dict[str, Any]: + if reference.shape != candidate.shape: + raise ValueError( + "Logit shape mismatch: " + f"reference={tuple(reference.shape)} candidate={tuple(candidate.shape)}." + ) + if not torch.isfinite(reference).all() or not torch.isfinite(candidate).all(): + raise ValueError("Logit artifacts must contain only finite values.") + diff = (reference - candidate).abs().flatten() + max_abs = float(diff.max().item()) + mean_abs = float(diff.mean().item()) + p99_abs = float(torch.quantile(diff, 0.99).item()) + reference_argmax = reference.argmax(dim=-1).tolist() + candidate_argmax = candidate.argmax(dim=-1).tolist() + gates = { + "max_abs_diff": max_abs <= float(max_abs_limit), + "mean_abs_diff": mean_abs <= float(mean_abs_limit), + "p99_abs_diff": p99_abs <= float(p99_abs_limit), + } + return { + "shape": list(reference.shape), + "max_abs_diff": max_abs, + "mean_abs_diff": mean_abs, + "p99_abs_diff": p99_abs, + "reference_argmax": reference_argmax, + "candidate_argmax": candidate_argmax, + "argmax_match": reference_argmax == candidate_argmax, + "limits": { + "max_abs_diff": float(max_abs_limit), + "mean_abs_diff": float(mean_abs_limit), + "p99_abs_diff": float(p99_abs_limit), + }, + "gates": gates, + "passed": all(gates.values()), + } + + +def _validate_generated_tokens( + reference_summary: dict[str, Any], + candidate_summary: dict[str, Any], + reference_prefill_logits: torch.Tensor, + *, + allow_first_token_tie: bool, + tie_atol: float, +) -> dict[str, Any]: + reference = [int(value) for value in reference_summary["generated_token_ids"]] + candidate = [int(value) for value in candidate_summary["generated_token_ids"]] + exact_match = reference == candidate + tie_accepted = False + reference_tie_gap = None + if ( + not exact_match + and allow_first_token_tie + and reference + and len(reference) == len(candidate) + and reference[1:] == candidate[1:] + ): + row = reference_prefill_logits.reshape( + -1, reference_prefill_logits.shape[-1] + )[-1] + candidate_first = candidate[0] + if not 0 <= candidate_first < int(row.numel()): + raise ValueError( + f"Candidate token {candidate_first} is outside the logit vocabulary." + ) + reference_tie_gap = float((row.max() - row[candidate_first]).item()) + tie_accepted = reference_tie_gap <= float(tie_atol) + return { + "reference": reference, + "candidate": candidate, + "exact_match": exact_match, + "first_token_tie_accepted": tie_accepted, + "reference_tie_gap": reference_tie_gap, + "tie_atol": float(tie_atol), + "passed": exact_match or tie_accepted, + } + + +def compare_runs( + reference_dir: Path, + candidate_dir: Path, + *, + max_abs_limit: float, + mean_abs_limit: float, + p99_abs_limit: float, + allow_first_token_tie: bool, + tie_atol: float, +) -> dict[str, Any]: + limits = (max_abs_limit, mean_abs_limit, p99_abs_limit, tie_atol) + if any(float(value) < 0.0 for value in limits): + raise ValueError(f"Comparison tolerances must be non-negative, got {limits}.") + reference_summary, reference_logits = _load_artifact(reference_dir) + candidate_summary, candidate_logits = _load_artifact(candidate_dir) + + mismatches = { + field: { + "reference": reference_summary.get(field), + "candidate": candidate_summary.get(field), + } + for field in MATCHED_RUN_FIELDS + if reference_summary.get(field) != candidate_summary.get(field) + } + if mismatches: + raise ValueError(f"Parallel runs are not comparable: {mismatches}.") + + comparisons = { + key: _compare_tensor( + reference_logits[key], + candidate_logits[key], + max_abs_limit=max_abs_limit, + mean_abs_limit=mean_abs_limit, + p99_abs_limit=p99_abs_limit, + ) + for key in LOGIT_KEYS + } + generated_tokens = _validate_generated_tokens( + reference_summary, + candidate_summary, + reference_logits["prefill"], + allow_first_token_tie=allow_first_token_tie, + tie_atol=tie_atol, + ) + gates = { + **{f"{key}_logits": value["passed"] for key, value in comparisons.items()}, + "generated_tokens": generated_tokens["passed"], + } + return { + "status": "success" if all(gates.values()) else "metric_failed", + "reference": { + "directory": str(reference_dir.resolve()), + "tensor_parallel_size": reference_summary["tensor_parallel_size"], + "expert_parallel_size": reference_summary["expert_parallel_size"], + "summary_sha256": _file_sha256(reference_dir / "summary.json"), + "logits_sha256": _file_sha256(reference_dir / "logits.pt"), + }, + "candidate": { + "directory": str(candidate_dir.resolve()), + "tensor_parallel_size": candidate_summary["tensor_parallel_size"], + "expert_parallel_size": candidate_summary["expert_parallel_size"], + "summary_sha256": _file_sha256(candidate_dir / "summary.json"), + "logits_sha256": _file_sha256(candidate_dir / "logits.pt"), + }, + "logits": comparisons, + "generated_tokens": generated_tokens, + "gates": gates, + } + + +def main() -> int: + args = _parse_args() + output_path = Path(args.output).expanduser().resolve() + output_path.parent.mkdir(parents=True, exist_ok=True) + try: + result = compare_runs( + Path(args.reference_dir).expanduser().resolve(), + Path(args.candidate_dir).expanduser().resolve(), + max_abs_limit=float(args.max_abs_diff), + mean_abs_limit=float(args.mean_abs_diff), + p99_abs_limit=float(args.p99_abs_diff), + allow_first_token_tie=bool(args.allow_first_token_tie), + tie_atol=float(args.tie_atol), + ) + except Exception as exc: + result = { + "status": "invalid_input", + "reference_dir": str( + Path(args.reference_dir).expanduser().resolve() + ), + "candidate_dir": str( + Path(args.candidate_dir).expanduser().resolve() + ), + "error_type": type(exc).__name__, + "error_message": str(exc), + "traceback": traceback.format_exc(), + } + output_path.write_text( + json.dumps(result, ensure_ascii=False, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + stream = sys.stdout if result["status"] == "success" else sys.stderr + print( + json.dumps(result, ensure_ascii=False, indent=2, sort_keys=True), + file=stream, + ) + return 0 if result["status"] == "success" else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/validation/validate_glm_checkpoint_parallel.py b/scripts/validation/validate_glm_checkpoint_parallel.py new file mode 100644 index 00000000..a94df294 --- /dev/null +++ b/scripts/validation/validate_glm_checkpoint_parallel.py @@ -0,0 +1,546 @@ +#!/usr/bin/env python3 +"""Validate one real GLM-4.7-Flash TP/EP runtime variant. + +This script deliberately uses the public ``LLM`` entry point so worker launch, +weight loading, scheduling, collectives, and cache allocation follow the same +path as serving. Cross-variant comparisons are performed separately from the +saved float32 logits artifacts. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import subprocess +import sys +import traceback +from datetime import datetime +from pathlib import Path +from typing import Any + +import torch + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Validate one real GLM checkpoint TP/EP variant." + ) + parser.add_argument("--model_path", required=True) + parser.add_argument("--output_dir", required=True) + parser.add_argument("--tensor_parallel_size", type=int, default=1) + parser.add_argument("--expert_parallel_size", type=int, default=1) + parser.add_argument( + "--sparse_method", + choices=( + "vanilla", + "streamingllm", + "snapkv", + "h2o", + "omnikv", + "rkv", + ), + default="vanilla", + ) + parser.add_argument( + "--decode_cuda_graph", + action=argparse.BooleanOptionalAction, + default=False, + ) + parser.add_argument("--seed", type=int, default=20260805) + parser.add_argument("--prompt_len", type=int, default=64) + parser.add_argument("--max_tokens", type=int, default=3) + parser.add_argument("--gpu_memory_utilization", type=float, default=0.80) + parser.add_argument("--mlp_chunk_size", type=int, default=64) + parser.add_argument("--max_steps", type=int, default=256) + return parser.parse_args() + + +def _write_json(path: Path, payload: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps(payload, ensure_ascii=False, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + + +def _git_output(*args: str) -> str: + result = subprocess.run( + ["git", *args], + check=False, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + return result.stdout.strip() + + +def _tensor_sha256(tensor: torch.Tensor) -> str: + value = tensor.detach().float().cpu().contiguous() + return hashlib.sha256(value.numpy().tobytes()).hexdigest() + + +def _tensor_evidence(tensor: torch.Tensor) -> dict[str, Any]: + value = tensor.detach().float().cpu().contiguous() + return { + "shape": [int(dim) for dim in value.shape], + "dtype": str(value.dtype), + "sha256": _tensor_sha256(value), + "argmax": value.argmax(dim=-1).tolist(), + "mean": float(value.mean().item()), + "max": float(value.max().item()), + "min": float(value.min().item()), + } + + +def _prompt_token_ids(*, seed: int, prompt_len: int, vocab_size: int) -> list[int]: + if prompt_len <= 0: + raise ValueError(f"prompt_len must be positive, got {prompt_len}.") + if vocab_size <= 1024: + raise ValueError( + "GLM checkpoint validation expects a vocabulary larger than 1024, " + f"got {vocab_size}." + ) + usable = vocab_size - 512 + base = 256 + int(seed) % max(1, usable - 256) + return [256 + ((base + 97 * index) % usable) for index in range(prompt_len)] + + +def _validate_parallel_evidence( + summaries: list[dict[str, Any]], + *, + tensor_parallel_size: int, + expert_parallel_size: int, + last_moe_layer: int, + require_all_ep_ranks_hit: bool = True, + require_decode_cuda_graph: bool = False, + require_decode_cuda_graph_execution: bool = False, +) -> dict[str, Any]: + from sparsevllm.distributed.parallel_context import ( + hybrid_moe_group_ranks, + parallel_group_ranks, + ) + + tp_size = int(tensor_parallel_size) + ep_size = int(expert_parallel_size) + hybrid_moe = tp_size > 1 and ep_size > 1 + expected_world_size = tp_size if hybrid_moe else tp_size * ep_size + if len(summaries) != expected_world_size: + raise RuntimeError( + "Expected one debug summary per world rank: " + f"expected={expected_world_size} got={len(summaries)}." + ) + + world_ranks = sorted(int(item["world_rank"]) for item in summaries) + if world_ranks != list(range(expected_world_size)): + raise RuntimeError(f"Unexpected world ranks: {world_ranks}.") + + if hybrid_moe: + raw_groups = hybrid_moe_group_ranks( + outer_tp_size=tp_size, + moe_ep_size=ep_size, + ) + group_ranks = { + "attention": raw_groups["attention"], + "expert": raw_groups["moe_expert"], + "moe_tensor": raw_groups["moe_tensor"], + "data": raw_groups["data"], + } + else: + raw_groups = parallel_group_ranks( + tp_size=tp_size, + ep_size=ep_size, + dp_size=1, + ) + group_ranks = { + "attention": raw_groups["tensor"], + "expert": raw_groups["expert"], + "moe_tensor": raw_groups["tensor"], + "data": raw_groups["data"], + } + + def expected_group(dimension: str, world_rank: int) -> dict[str, Any]: + ranks = next( + ranks + for ranks in group_ranks[dimension] + if int(world_rank) in ranks + ) + return { + "rank": ranks.index(int(world_rank)), + "size": len(ranks), + "ranks": list(ranks), + } + + local_ranges: list[tuple[int, int]] = [] + local_hit_counts: list[int] = [] + replica_checks: list[dict[str, Any]] = [] + layer_key = str(int(last_moe_layer)) + for summary in summaries: + configured = summary["parallel"]["configured"] + expected_configured = { + "tensor_parallel_size": int(tensor_parallel_size), + "expert_parallel_size": int(expert_parallel_size), + "data_parallel_size": 1, + "world_size": expected_world_size, + } + if configured != expected_configured: + raise RuntimeError( + "Configured parallel evidence mismatch: " + f"expected={expected_configured} got={configured}." + ) + effective = summary["parallel"]["effective"] + world = effective["world"] + if int(world["size"]) != expected_world_size: + raise RuntimeError(f"Effective world-size mismatch: {world}.") + if world != { + "rank": int(summary["world_rank"]), + "size": expected_world_size, + "ranks": list(range(expected_world_size)), + }: + raise RuntimeError( + "Effective world group mismatch: " + f"world_rank={summary['world_rank']} group={world}." + ) + for dimension in ("attention", "expert", "moe_tensor", "data"): + expected = expected_group(dimension, int(summary["world_rank"])) + if effective[dimension] != expected: + raise RuntimeError( + f"Effective {dimension} group mismatch on world rank " + f"{summary['world_rank']}: expected={expected} " + f"got={effective[dimension]}." + ) + + layer = summary.get("moe_local", {}).get(layer_key) + if layer is None: + raise RuntimeError( + f"Missing MoE debug evidence for layer={last_moe_layer} " + f"on world_rank={summary['world_rank']}." + ) + local_ranges.append( + (int(layer["local_expert_start"]), int(layer["local_expert_end"])) + ) + local_hit_counts.append(int(layer["local_hit_count"])) + + consistency = summary.get("replica_consistency") + if not isinstance(consistency, dict): + raise RuntimeError( + "Missing replica-consistency evidence for " + f"world_rank={summary['world_rank']}." + ) + replica_checks.append(consistency) + logits_tolerance_ratio = consistency["last_logits_tolerance_ratio"] + if tensor_parallel_size > 1: + if ( + logits_tolerance_ratio is not None + or consistency.get("last_logits_comparison") + != "not_applicable_tp_vocab_sharded" + ): + raise RuntimeError( + "TP-sharded logits must be marked not applicable for cross-rank " + f"comparison: world_rank={summary['world_rank']} " + f"evidence={consistency}." + ) + elif logits_tolerance_ratio is None or float(logits_tolerance_ratio) > 1.0: + raise RuntimeError( + "Cross-rank logits exceeded the declared tolerance or were not " + f"compared: world_rank={summary['world_rank']} evidence={consistency}." + ) + moe_consistency = consistency.get("moe_layers", {}).get(layer_key) + if moe_consistency is None: + raise RuntimeError( + f"Missing cross-rank MoE consistency for layer={layer_key}." + ) + if bool(moe_consistency["topk_ids_mismatch"]): + raise RuntimeError( + f"Router top-k differs across ranks: {moe_consistency}." + ) + for field in ("topk_weights_tolerance_ratio", "output_tolerance_ratio"): + if float(moe_consistency[field]) > 1.0: + raise RuntimeError( + f"Cross-rank MoE {field} exceeded tolerance: " + f"{moe_consistency}." + ) + graph = summary.get("decode_cuda_graph") + if not isinstance(graph, dict): + raise RuntimeError( + "Missing decode CUDA Graph evidence on world rank " + f"{summary['world_rank']}." + ) + if bool(graph.get("enabled")) != bool(require_decode_cuda_graph): + raise RuntimeError( + "Decode CUDA Graph config evidence mismatch on world rank " + f"{summary['world_rank']}: {graph}." + ) + if require_decode_cuda_graph_execution: + if not require_decode_cuda_graph: + raise ValueError( + "Requiring decode CUDA Graph execution also requires the " + "graph config to be enabled." + ) + if int(graph.get("capture_count", 0)) <= 0: + raise RuntimeError( + f"World rank {summary['world_rank']} did not capture a decode graph." + ) + if int(graph.get("replay_count", 0)) <= 0: + raise RuntimeError( + f"World rank {summary['world_rank']} did not replay a decode graph." + ) + if int(graph.get("eager_static_count", 0)) != 0 or int( + graph.get("force_eager_count", 0) + ) != 0: + raise RuntimeError( + "Decode CUDA Graph used an eager path on world rank " + f"{summary['world_rank']}: {graph}." + ) + + if ep_size > 1: + expected_width = 64 // ep_size + expected_ranges = [ + (rank * expected_width, (rank + 1) * expected_width) + for rank in range(ep_size) + for _ in range(tp_size // ep_size if hybrid_moe else tp_size) + ] + if sorted(local_ranges) != expected_ranges: + raise RuntimeError( + "EP local expert shards do not partition 64 experts: " + f"expected={expected_ranges} got={sorted(local_ranges)}." + ) + hits_by_range: dict[tuple[int, int], int] = {} + for local_range, hit_count in zip(local_ranges, local_hit_counts): + hits_by_range[local_range] = ( + hits_by_range.get(local_range, 0) + int(hit_count) + ) + if require_all_ep_ranks_hit and any( + hit_count <= 0 for hit_count in hits_by_range.values() + ): + raise RuntimeError( + "At least one EP rank did not execute a local routed expert: " + f"hits_by_range={hits_by_range}." + ) + expected_replicated_attention = tp_size == 1 + if any( + bool(summary["parallel"]["attention_replicated_for_ep"]) + != expected_replicated_attention + for summary in summaries + ): + raise RuntimeError( + "EP replicated-attention evidence does not match the selected " + f"layout: TP={tp_size}, EP={ep_size}." + ) + + return { + "world_ranks": world_ranks, + "local_expert_ranges": [list(value) for value in local_ranges], + "local_expert_hit_counts": local_hit_counts, + "all_ep_ranks_hit": all(hit_count > 0 for hit_count in local_hit_counts), + "hybrid_moe": hybrid_moe, + "expected_groups": { + key: [list(ranks) for ranks in value] + for key, value in group_ranks.items() + }, + "replica_checks": replica_checks, + } + + +def _run(args: argparse.Namespace, output_dir: Path) -> dict[str, Any]: + os.environ["SPARSEVLLM_DEBUG_RUNTIME"] = "1" + os.environ["SPARSEVLLM_DEBUG_MOE"] = "1" + + from sparsevllm import LLM, SamplingParams + + tp_size = int(args.tensor_parallel_size) + ep_size = int(args.expert_parallel_size) + if tp_size <= 0 or ep_size <= 0: + raise ValueError(f"TP and EP must be positive, got TP={tp_size}, EP={ep_size}.") + if tp_size > 1 and ep_size > 1 and tp_size % ep_size: + raise ValueError( + "Joint GLM outer TP and MoE EP requires TP divisible by EP, got " + f"TP={tp_size}, EP={ep_size}." + ) + expected_devices = ( + tp_size if tp_size > 1 and ep_size > 1 else tp_size * ep_size + ) + actual_devices = torch.cuda.device_count() + if actual_devices != expected_devices: + raise RuntimeError( + "CUDA_VISIBLE_DEVICES must expose exactly the requested world size: " + f"expected={expected_devices} got={actual_devices}." + ) + + torch.manual_seed(int(args.seed)) + max_model_len = int(args.prompt_len) + int(args.max_tokens) + 16 + llm = None + try: + llm = LLM( + args.model_path, + tensor_parallel_size=tp_size, + expert_parallel_size=ep_size, + data_parallel_size=1, + sparse_method=args.sparse_method, + enforce_eager=True, + decode_cuda_graph=bool(args.decode_cuda_graph), + max_model_len=max_model_len, + max_num_seqs_in_batch=1, + max_decoding_seqs=1, + max_num_batched_tokens=max(int(args.prompt_len), 1), + engine_prefill_chunk_size=max(int(args.prompt_len), 1), + mlp_chunk_size=int(args.mlp_chunk_size), + gpu_memory_utilization=float(args.gpu_memory_utilization), + throughput_log_interval_s=0.0, + weight_loading_workers=1, + ) + prompt = _prompt_token_ids( + seed=int(args.seed), + prompt_len=int(args.prompt_len), + vocab_size=int(llm.config.hf_config.vocab_size), + ) + sampling_params = SamplingParams( + temperature=0.0, + top_p=1.0, + ignore_eos=True, + max_tokens=int(args.max_tokens), + ) + llm.add_request(prompt, sampling_params) + + prefill_logits = None + prefill_rank_summaries = None + finished_outputs: list[tuple[int, list[int], Any, Any]] = [] + steps = 0 + while not llm.is_finished(): + step_outputs, num_tokens = llm.step() + steps += 1 + if steps > int(args.max_steps): + raise RuntimeError( + f"Validation exceeded max_steps={int(args.max_steps)}." + ) + if ( + num_tokens > 0 + and llm.last_step_token_outputs + and prefill_rank_summaries is None + ): + prefill_logits = llm.debug_last_logits().detach().float().cpu() + prefill_rank_summaries = llm.debug_sparse_state_summaries() + finished_outputs.extend(step_outputs) + + if prefill_logits is None or prefill_rank_summaries is None: + raise RuntimeError("The run completed without final-prefill evidence.") + if len(finished_outputs) != 1: + raise RuntimeError( + f"Expected one finished request, got {len(finished_outputs)}." + ) + generated_ids = [int(value) for value in finished_outputs[0][1]] + if len(generated_ids) != int(args.max_tokens): + raise RuntimeError( + "Greedy generation returned an unexpected token count: " + f"expected={int(args.max_tokens)} got={len(generated_ids)}." + ) + final_logits = llm.debug_last_logits().detach().float().cpu() + final_rank_summaries = llm.debug_sparse_state_summaries() + last_moe_layer = int(llm.config.hf_config.num_hidden_layers) - 1 + + prefill_parallel = _validate_parallel_evidence( + prefill_rank_summaries, + tensor_parallel_size=tp_size, + expert_parallel_size=ep_size, + last_moe_layer=last_moe_layer, + require_decode_cuda_graph=bool(args.decode_cuda_graph), + ) + final_parallel = _validate_parallel_evidence( + final_rank_summaries, + tensor_parallel_size=tp_size, + expert_parallel_size=ep_size, + last_moe_layer=last_moe_layer, + # A single decode token selects only top-k experts and does not + # have to cover every EP shard. Prefill coverage above proves all + # local expert partitions execute; keep the final per-rank counts + # as evidence without imposing an invalid per-token requirement. + require_all_ep_ranks_hit=False, + require_decode_cuda_graph=bool(args.decode_cuda_graph), + require_decode_cuda_graph_execution=bool(args.decode_cuda_graph), + ) + + torch.save( + { + "prefill": prefill_logits, + "final_decode": final_logits, + }, + output_dir / "logits.pt", + ) + _write_json(output_dir / "prefill_rank_summaries.json", prefill_rank_summaries) + _write_json(output_dir / "final_rank_summaries.json", final_rank_summaries) + + return { + "status": "success", + "real_checkpoint": True, + "model_path": str(Path(args.model_path).resolve()), + "tensor_parallel_size": tp_size, + "expert_parallel_size": ep_size, + "data_parallel_size": 1, + "world_size": expected_devices, + "sparse_method": str(args.sparse_method), + "decode_cuda_graph": bool(args.decode_cuda_graph), + "seed": int(args.seed), + "prompt_len": len(prompt), + "prompt_sha256": hashlib.sha256( + torch.tensor(prompt, dtype=torch.int64).numpy().tobytes() + ).hexdigest(), + "max_tokens": int(args.max_tokens), + "generated_token_ids": generated_ids, + "steps": steps, + "prefill_logits": _tensor_evidence(prefill_logits), + "final_decode_logits": _tensor_evidence(final_logits), + "prefill_parallel": prefill_parallel, + "final_parallel": final_parallel, + "logits_artifact": str(output_dir / "logits.pt"), + } + finally: + if llm is not None: + llm.exit() + + +def main() -> int: + args = _parse_args() + output_dir = Path(args.output_dir).expanduser().resolve() + output_dir.mkdir(parents=True, exist_ok=True) + run_info = { + "status": "running", + "started_at": datetime.now().astimezone().isoformat(timespec="seconds"), + "command": [sys.executable, *sys.argv], + "cwd": os.getcwd(), + "git_commit": _git_output("rev-parse", "HEAD"), + "git_status": _git_output("status", "--short"), + "cuda_visible_devices": os.getenv("CUDA_VISIBLE_DEVICES"), + "args": vars(args), + } + _write_json(output_dir / "run_info.json", run_info) + try: + summary = _run(args, output_dir) + except BaseException as exc: + failure = { + "status": "model_failed", + "error_type": type(exc).__name__, + "error_message": str(exc), + "traceback": traceback.format_exc(), + } + _write_json(output_dir / "summary.json", failure) + run_info.update( + status="model_failed", + finished_at=datetime.now().astimezone().isoformat(timespec="seconds"), + error=failure, + ) + _write_json(output_dir / "run_info.json", run_info) + raise + _write_json(output_dir / "summary.json", summary) + run_info.update( + status="success", + finished_at=datetime.now().astimezone().isoformat(timespec="seconds"), + summary_path=str(output_dir / "summary.json"), + ) + _write_json(output_dir / "run_info.json", run_info) + print(json.dumps(summary, ensure_ascii=False, indent=2, sort_keys=True)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/sparsevllm/configs/model.py b/src/sparsevllm/configs/model.py index 635ac201..2cb3503e 100644 --- a/src/sparsevllm/configs/model.py +++ b/src/sparsevllm/configs/model.py @@ -9,7 +9,6 @@ from transformers import AutoConfig from sparsevllm.method_registry import ( - normalize_sparse_method, validate_model_runtime_compatibility, validate_sparse_method_assets, ) @@ -25,6 +24,30 @@ def _config_to_namespace(config: dict[str, Any]) -> SimpleNamespace: return SimpleNamespace(**config) +def resolve_attention_qk_head_dim(hf_config: Any) -> int: + qk_nope_head_dim = config_get(hf_config, "qk_nope_head_dim", None) + qk_rope_head_dim = config_get(hf_config, "qk_rope_head_dim", None) + if qk_nope_head_dim is not None and qk_rope_head_dim is not None: + head_dim = int(qk_nope_head_dim) + int(qk_rope_head_dim) + if head_dim <= 0: + raise ValueError(f"Attention QK head dimension must be positive, got {head_dim}.") + return head_dim + head_dim = config_get(hf_config, "head_dim", None) + if head_dim is not None: + head_dim = int(head_dim) + if head_dim <= 0: + raise ValueError(f"Attention QK head dimension must be positive, got {head_dim}.") + return head_dim + hidden_size = int(config_get(hf_config, "hidden_size", 0) or 0) + num_heads = int(config_get(hf_config, "num_attention_heads", 0) or 0) + if hidden_size <= 0 or num_heads <= 0 or hidden_size % num_heads: + raise ValueError( + "Attention QK head dimension requires valid head_dim or divisible " + "hidden_size/num_attention_heads." + ) + return hidden_size // num_heads + + def _load_model_config(model_path: str) -> Any: try: return AutoConfig.from_pretrained(model_path, trust_remote_code=True) @@ -70,17 +93,9 @@ def _resolve_attention_cache_layout(config, model_type: str) -> None: ) if config.attention_cache_layout == "explicit_kv": return - if normalize_sparse_method(config.vllm_sparse_method): - raise NotImplementedError( - "GLM-4.7-Flash latent MLA supports only vanilla attention." - ) - if config.enable_prefix_caching or config.enable_prefix_cache_offload: - raise NotImplementedError( - "GLM-4.7-Flash latent MLA does not support prefix caching or offload." - ) - if config.decode_cuda_graph or not config.enforce_eager: + if config.enable_prefix_cache_offload: raise NotImplementedError( - "GLM-4.7-Flash latent MLA requires eager execution." + "GLM-4.7-Flash latent MLA does not support prefix cache offload." ) if config.quantization_config.enabled: raise NotImplementedError( @@ -90,13 +105,23 @@ def _resolve_attention_cache_layout(config, model_type: str) -> None: raise NotImplementedError( "GLM-4.7-Flash latent MLA requires BF16 checkpoints." ) - if ( - int(config_get(config.hf_config, "kv_lora_rank", 0)) != 512 - or int(config_get(config.hf_config, "qk_rope_head_dim", 0)) != 64 - ): + expected = { + "num_attention_heads": 20, + "q_lora_rank": 768, + "kv_lora_rank": 512, + "qk_nope_head_dim": 192, + "qk_rope_head_dim": 64, + "v_head_dim": 256, + } + invalid = { + field: config_get(config.hf_config, field, None) + for field, value in expected.items() + if config_get(config.hf_config, field, None) != value + } + if invalid: + field, actual = next(iter(invalid.items())) raise NotImplementedError( - "GLM-4.7-Flash latent MLA requires kv_lora_rank=512 and " - "qk_rope_head_dim=64." + f"GLM-4.7-Flash latent MLA requires {field}={expected[field]}, got {actual}." ) diff --git a/src/sparsevllm/configs/prefix_cache.py b/src/sparsevllm/configs/prefix_cache.py index db252397..5970e008 100644 --- a/src/sparsevllm/configs/prefix_cache.py +++ b/src/sparsevllm/configs/prefix_cache.py @@ -100,8 +100,8 @@ def normalize_prefix_cache(config) -> None: config.recurrent_state_max_bytes = recurrent_state_max_bytes if config.enable_prefix_caching and config.vllm_sparse_method not in PREFIX_CACHE_SUPPORTED_METHODS: raise ValueError( - "prefix caching only supports vanilla, omnikv, quest, snapkv, " - "h2o, pyramidkv, rkv, and skipkv." + "prefix caching only supports vanilla, streamingllm, omnikv, quest, " + "snapkv, h2o, pyramidkv, rkv, and skipkv." ) config.prefix_cache_salt = str(config.prefix_cache_salt or "") diff --git a/src/sparsevllm/engine/cache_manager/__init__.py b/src/sparsevllm/engine/cache_manager/__init__.py index c27f80fb..b80dc107 100644 --- a/src/sparsevllm/engine/cache_manager/__init__.py +++ b/src/sparsevllm/engine/cache_manager/__init__.py @@ -2,6 +2,7 @@ from .base import ( AttentionCacheWrite, + AttentionKeyComputeView, AttentionPayload, AttentionViewMeta, CacheManager, @@ -17,6 +18,7 @@ __all__ = [ "AttentionCacheWrite", + "AttentionKeyComputeView", "AttentionPayload", "AttentionViewMeta", "CacheManager", diff --git a/src/sparsevllm/engine/cache_manager/base.py b/src/sparsevllm/engine/cache_manager/base.py index 54e92b3f..7dba1f86 100644 --- a/src/sparsevllm/engine/cache_manager/base.py +++ b/src/sparsevllm/engine/cache_manager/base.py @@ -5,7 +5,7 @@ from collections import deque from dataclasses import dataclass, fields, is_dataclass from abc import ABC, abstractmethod -from typing import Any +from typing import Any, Callable import torch import torch.distributed as dist @@ -20,6 +20,7 @@ from sparsevllm.method_registry import SUPPORTED_SPARSE_METHODS, normalize_sparse_method from sparsevllm.triton_kernel.store_kvcache import store_kvcache import sparsevllm.platforms as platforms +from sparsevllm.configs.model import resolve_attention_qk_head_dim from sparsevllm.utils.log import logger, log_level @@ -209,6 +210,17 @@ class PrefillComputeView: payload: AttentionPayload +@dataclass(frozen=True) +class AttentionKeyComputeView: + """Physical payload plus the slots whose actual attention keys are needed.""" + + active_slots: torch.Tensor + payload: AttentionPayload + + +AttentionKeyMaterializer = Callable[[AttentionKeyComputeView], torch.Tensor] + + class CacheManager(ABC): """每个 Rank 只有一个 CacheManager，内部管理所有层的物理槽位和 KV Cache。""" @@ -233,11 +245,7 @@ def __init__(self, config: Config, parallel_context: ParallelContext): self.num_kv_layers = int(self.runtime_layout.num_kv_layers) self.num_kv_heads = self.hf_config.num_key_value_heads // self.tp_size - self.head_dim = getattr( - self.hf_config, - "head_dim", - self.hf_config.hidden_size // self.hf_config.num_attention_heads, - ) + self.head_dim = resolve_attention_qk_head_dim(self.hf_config) self.max_model_len = config.max_model_len resident_buffer_rows = int(config.max_num_seqs_in_gpu) @@ -264,6 +272,7 @@ def __init__(self, config: Config, parallel_context: ParallelContext): self.kv_cache = None self._decode_static_max_context_len: int | None = None self._raw_offload_prefill_phases: dict[int, bool] = {} + self._attention_key_materializers: dict[int, AttentionKeyMaterializer] = {} def synchronize_prefix_cache_delete_plan( self, @@ -708,6 +717,108 @@ def get_layer_compute_payload( context_lens, ) + def register_attention_key_materializer( + self, + layer_idx: int, + materializer: AttentionKeyMaterializer, + ) -> None: + """Bind a model/operator hook that reconstructs actual keys from a layout.""" + + layer_idx = int(layer_idx) + self.kv_layer_index(layer_idx) + if not callable(materializer): + raise TypeError( + "Attention key materializer must be callable, got " + f"{type(materializer).__name__}." + ) + registry = getattr(self, "_attention_key_materializers", None) + if registry is None: + registry = {} + self._attention_key_materializers = registry + existing = registry.get(layer_idx) + if existing is not None and existing != materializer: + raise RuntimeError( + "Attention key materializer is already bound for " + f"layer={layer_idx}." + ) + registry[layer_idx] = materializer + + def has_attention_key_materializer(self, layer_idx: int) -> bool: + registry = getattr(self, "_attention_key_materializers", {}) + return int(layer_idx) in registry + + def build_attention_key_compute_view( + self, + layer_idx: int, + active_slots: torch.Tensor, + ) -> AttentionKeyComputeView: + """Build a tagged view without assuming an explicit or latent layout.""" + + layer_idx = int(layer_idx) + kv_idx = self.kv_layer_index(layer_idx) + storage = getattr(self, "attention_cache_storage", None) + if storage is None: + k_cache, v_cache = self.get_layer_kv_cache(layer_idx) + payload: AttentionPayload = ExplicitKVPayload( + k_cache=k_cache, + v_cache=v_cache, + ) + else: + payload = storage.layer_payload(kv_idx) + return AttentionKeyComputeView( + active_slots=active_slots, + payload=payload, + ) + + @torch.no_grad() + def materialize_attention_keys( + self, + layer_idx: int, + active_slots: torch.Tensor, + ) -> torch.Tensor: + """Return the actual post-RoPE per-head keys for arbitrary cache slots.""" + + view = self.build_attention_key_compute_view(layer_idx, active_slots) + slots = view.active_slots + if slots.ndim == 0: + raise ValueError("Attention key slots must have at least one dimension.") + if slots.dtype not in (torch.int32, torch.int64): + raise TypeError( + "Attention key slots must use int32 or int64, got " + f"{slots.dtype}." + ) + + if isinstance(view.payload, ExplicitKVPayload): + k_cache = view.payload.k_cache + flat_slots = slots.to(device=k_cache.device, dtype=torch.long).reshape(-1) + keys = k_cache.index_select(0, flat_slots).view( + *slots.shape, + *k_cache.shape[1:], + ) + else: + registry = getattr(self, "_attention_key_materializers", {}) + materializer = registry.get(int(layer_idx)) + if materializer is None: + raise RuntimeError( + "The attention cache layout requires an actual-key " + f"materializer at layer={int(layer_idx)}." + ) + keys = materializer(view) + + expected_ndim = int(slots.ndim) + 2 + if keys.ndim != expected_ndim or tuple(keys.shape[: slots.ndim]) != tuple(slots.shape): + raise RuntimeError( + "Attention key materializer returned an invalid shape: " + f"slots={tuple(slots.shape)} keys={tuple(keys.shape)}; expected " + "the slot shape followed by [heads, head_dim]." + ) + if keys.device != slots.device: + raise RuntimeError( + "Materialized attention keys must share the slot device: " + f"slots={slots.device} keys={keys.device}." + ) + return keys + def get_prefill_compute_view( self, layer_idx: int, @@ -887,6 +998,31 @@ def decode_cuda_graph_keepalive_tensors(self) -> list[torch.Tensor]: """Cache-manager-owned tensors captured by decode CUDA graphs.""" return [] + def validate_decode_cuda_graph_slot_mappings(self) -> None: + """Validate every layer's static decode mapping as one store scope. + + Static preparation calls this for every decode step. Graph capture then + calls it again after its eager warmup because a storage backend may + consume the one-forward prevalidation during that warmup. Sparse + methods may bind a different stable mapping tensor per layer, so the + storage is given the exact layer-ordered set instead of only the + runner's common metadata buffer. + """ + storage = getattr(self, "attention_cache_storage", None) + if storage is None: + return + slot_mappings: list[torch.Tensor] = [] + for layer_idx in self.kv_transformer_layer_indices(): + state = self.get_layer_batch_states(layer_idx) + slot_mapping = state.slot_mapping + if slot_mapping is None: + raise RuntimeError( + "Decode CUDA graph capture requires a slot mapping for " + f"layer={layer_idx}." + ) + slot_mappings.append(slot_mapping) + storage.validate_slot_mappings(tuple(slot_mappings)) + def decode_cuda_graph_max_cached_graphs(self) -> int | None: """Optional bound for captured decode graph states. @@ -1426,9 +1562,7 @@ def _dense_baseline_bytes(self) -> int: f"heads={global_heads} tp={attention_tp_size}." ) local_heads = global_heads // attention_tp_size - qk_head_dim = int(hf_config.qk_nope_head_dim) + int( - hf_config.qk_rope_head_dim - ) + qk_head_dim = resolve_attention_qk_head_dim(hf_config) value_head_dim = int(hf_config.v_head_dim) values_per_token = local_heads * (qk_head_dim + value_head_dim) else: diff --git a/src/sparsevllm/engine/cache_manager/h2o.py b/src/sparsevllm/engine/cache_manager/h2o.py index e935f249..85160f60 100644 --- a/src/sparsevllm/engine/cache_manager/h2o.py +++ b/src/sparsevllm/engine/cache_manager/h2o.py @@ -13,6 +13,7 @@ from .base import ExplicitKVPayload, PrefillComputeView from .snapkv import SnapKVCacheManager +from .storage import ExplicitKVStorage class H2OCacheManager(SnapKVCacheManager): @@ -417,6 +418,8 @@ def prepare_decode_static( state.req_indices = layers_req_indices[layer_id] self._decode_static_state_binding_key = binding_key + self.validate_decode_cuda_graph_slot_mappings() + slot_mapping.copy_(layers_slot_mapping[first_layer]) context_lens.copy_(layers_context_lens[first_layer]) req_indices.copy_(layers_req_indices[first_layer]) @@ -1124,7 +1127,7 @@ def _compact_final_prefill_dense_batch( """Move final H2O selections into ascending physical destination slots.""" if not seqs: raise RuntimeError("H2O final-prefill dense compaction requires sequences.") - self.kv_layer_index(layer_idx) + kv_idx = self.kv_layer_index(layer_idx) budget = self.h2o_decode_budget batch_size = len(seqs) keep_indices = keep_indices.to( @@ -1186,13 +1189,19 @@ def _compact_final_prefill_dense_batch( f"capacity={int(free_stack.numel())}." ) - k_cache, v_cache = self.get_layer_kv_cache(layer_idx) - workspace = self._get_final_prefill_workspace( - batch_size=batch_size, - budget=budget, - k_cache=k_cache, - v_cache=v_cache, - ) + storage = getattr(self, "attention_cache_storage", None) + uses_explicit_kv = storage is None or isinstance(storage, ExplicitKVStorage) + if uses_explicit_kv: + k_cache, v_cache = self.get_layer_kv_cache(layer_idx) + workspace = self._get_final_prefill_workspace( + batch_size=batch_size, + budget=budget, + k_cache=k_cache, + v_cache=v_cache, + ) + slot_capacity = int(k_cache.shape[0]) + else: + slot_capacity = storage.slot_capacity() rows_gpu = torch.tensor(row_indices, dtype=torch.long, device=self.device) old_slots = self.buffer_req_to_token_slots[layer_idx][ rows_gpu, :kv_len @@ -1209,9 +1218,9 @@ def _compact_final_prefill_dense_batch( f"layer={layer_idx}.", ) self._assert_final_prefill_tensor( - ((old_slots >= 0) & (old_slots < int(k_cache.shape[0]))).all(), + ((old_slots >= 0) & (old_slots < slot_capacity)).all(), "H2O final-prefill slot map contains an out-of-range physical slot: " - f"layer={layer_idx} num_slots={int(k_cache.shape[0])}.", + f"layer={layer_idx} num_slots={slot_capacity}.", ) globally_sorted_slots = torch.sort(old_slots.reshape(-1)).values @@ -1228,42 +1237,44 @@ def _compact_final_prefill_dense_batch( released_slots = sorted_old_slots[:, budget:].reshape(-1).contiguous() selected_flat = selected_slots.reshape(-1) - workspace[0].copy_( - k_cache.index_select(0, selected_flat).view( - batch_size, - budget, - int(k_cache.shape[1]), - int(k_cache.shape[2]), + destination_flat = destination_slots.reshape(-1).to(torch.long) + if uses_explicit_kv: + workspace[0].copy_( + k_cache.index_select(0, selected_flat).view( + batch_size, + budget, + int(k_cache.shape[1]), + int(k_cache.shape[2]), + ) ) - ) - workspace[1].copy_( - v_cache.index_select(0, selected_flat).view( - batch_size, - budget, - int(v_cache.shape[1]), - int(v_cache.shape[2]), + workspace[1].copy_( + v_cache.index_select(0, selected_flat).view( + batch_size, + budget, + int(v_cache.shape[1]), + int(v_cache.shape[2]), + ) ) - ) - - destination_flat = destination_slots.reshape(-1).to(torch.long) - k_cache.index_copy_( - 0, - destination_flat, - workspace[0].reshape( - batch_size * budget, - int(k_cache.shape[1]), - int(k_cache.shape[2]), - ), - ) - v_cache.index_copy_( - 0, - destination_flat, - workspace[1].reshape( - batch_size * budget, - int(v_cache.shape[1]), - int(v_cache.shape[2]), - ), - ) + k_cache.index_copy_( + 0, + destination_flat, + workspace[0].reshape( + batch_size * budget, + int(k_cache.shape[1]), + int(k_cache.shape[2]), + ), + ) + v_cache.index_copy_( + 0, + destination_flat, + workspace[1].reshape( + batch_size * budget, + int(v_cache.shape[1]), + int(v_cache.shape[2]), + ), + ) + else: + storage.copy_slots(kv_idx, selected_flat, destination_flat) free_stack[free_ptr : free_ptr + free_count] = released_slots.to( dtype=free_stack.dtype, diff --git a/src/sparsevllm/engine/cache_manager/rkv.py b/src/sparsevllm/engine/cache_manager/rkv.py index 378496ed..3746e81e 100644 --- a/src/sparsevllm/engine/cache_manager/rkv.py +++ b/src/sparsevllm/engine/cache_manager/rkv.py @@ -154,6 +154,69 @@ def _clear_rkv_query_cache_rows(self, layer_idx: int, row_indices: list[int | No rows_tensor = torch.tensor(rows, dtype=torch.long, device=self.device) positions[rows_tensor].fill_(-1) + @staticmethod + def attention_scores_from_materialized_keys( + q_window: torch.Tensor, + keys: torch.Tensor, + query_positions: torch.Tensor, + candidate_positions: torch.Tensor, + ) -> torch.Tensor: + """Match the R-KV prefill-score contract using actual post-RoPE keys.""" + + if q_window.ndim != 3 or keys.ndim != 3: + raise ValueError( + "R-KV materialized scoring requires q=[queries, heads, dim] " + f"and keys=[tokens, kv_heads, dim], got {tuple(q_window.shape)} " + f"and {tuple(keys.shape)}." + ) + if int(q_window.shape[-1]) != int(keys.shape[-1]): + raise ValueError( + "R-KV query/key head dimensions differ: " + f"q={int(q_window.shape[-1])} k={int(keys.shape[-1])}." + ) + num_query_heads = int(q_window.shape[1]) + num_key_heads = int(keys.shape[1]) + if num_key_heads <= 0 or num_query_heads % num_key_heads: + raise ValueError( + "R-KV query heads must be divisible by materialized key heads: " + f"q_heads={num_query_heads} key_heads={num_key_heads}." + ) + if query_positions.shape != (int(q_window.shape[0]),): + raise ValueError( + "R-KV query positions do not match the query window: " + f"positions={tuple(query_positions.shape)} " + f"queries={int(q_window.shape[0])}." + ) + if candidate_positions.shape != (int(keys.shape[0]),): + raise ValueError( + "R-KV candidate positions do not match materialized keys: " + f"positions={tuple(candidate_positions.shape)} " + f"keys={int(keys.shape[0])}." + ) + if int(keys.shape[0]) == 0: + return torch.empty((0,), dtype=torch.float32, device=keys.device) + + heads_per_key = num_query_heads // num_key_heads + expanded_keys = keys.repeat_interleave(heads_per_key, dim=1) + logits = torch.einsum( + "qhd,chd->qhc", + q_window.float(), + expanded_keys.float(), + ) + logits.mul_(float(q_window.shape[-1]) ** -0.5) + valid = query_positions[:, None] >= candidate_positions[None, :] + valid_rows = valid.any(dim=1) + logits = logits.masked_fill(~valid[:, None, :], float("-inf")) + logits = torch.where( + valid_rows[:, None, None], + logits, + torch.zeros_like(logits), + ) + probabilities = torch.softmax(logits, dim=-1) + probabilities.masked_fill_(~valid[:, None, :], 0.0) + probabilities.masked_fill_(~valid_rows[:, None, None], 0.0) + return probabilities.mean(dim=0).amax(dim=0) + def free_seq(self, seq_id: int): row_by_layer = [ self.seq_id_to_row[layer_idx].get(int(seq_id)) @@ -372,6 +435,44 @@ def rkv_query_attention_scores( ) q_window = cache[row_idx, cols].contiguous() + if self.has_attention_key_materializer(layer_idx): + candidate_start = max(0, int(candidate_start)) + candidate_end = max( + candidate_start, + kv_len - int(num_recent_tokens), + ) + scores = torch.zeros( + (kv_len,), + dtype=self._prefill_score_dtype(), + device=self.device, + ) + if candidate_end <= candidate_start: + return scores + candidate_positions = torch.arange( + candidate_start, + candidate_end, + dtype=torch.long, + device=self.device, + ) + candidate_slots = self.buffer_req_to_token_slots[layer_idx][ + row_idx, + candidate_start:candidate_end, + ] + keys = self.materialize_attention_keys( + layer_idx, + candidate_slots, + ) + candidate_scores = self.attention_scores_from_materialized_keys( + q_window, + keys, + positions, + candidate_positions, + ) + scores[candidate_start:candidate_end] = candidate_scores.to( + scores.dtype + ) + return scores + k_cache, _ = self.get_layer_kv_cache(layer_idx) attn_score = torch.zeros( (1, kv_len), @@ -426,6 +527,23 @@ def rkv_query_attention_scores_batch( "rkv_query_attention_scores_batch expected one kv_len per sequence: " f"seqs={len(seqs)} kv_lens={len(kv_lens)}" ) + if self.has_attention_key_materializer(layer_idx): + max_kv_len = max(int(kv_len) for kv_len in kv_lens) + scores = torch.zeros( + (len(seqs), max_kv_len), + dtype=self._prefill_score_dtype(), + device=self.device, + ) + for batch_idx, (seq, kv_len) in enumerate(zip(seqs, kv_lens)): + single = self.rkv_query_attention_scores( + layer_idx, + seq, + int(kv_len), + candidate_start=candidate_start, + num_recent_tokens=num_recent_tokens, + ) + scores[batch_idx, : int(kv_len)] = single + return scores obs = int(self._rkv_observation_tokens) if obs <= 0: @@ -624,9 +742,11 @@ def select_rkv_indices( configured_window = int(self.config.rkv_redundancy_window) window = int(slots.numel()) if configured_window == 0 else min(configured_window, int(slots.numel())) if window > 0: - k_cache, _ = self.get_layer_kv_cache(layer_idx) window_slots = slots[-window:] - window_keys = k_cache.index_select(0, window_slots) + window_keys = self.materialize_attention_keys( + layer_idx, + window_slots, + ) redundancy = self.redundancy_scores_from_keys( window_keys, similarity_threshold=float(self.config.rkv_similarity_threshold), @@ -697,13 +817,10 @@ def select_rkv_indices_batch( configured_window = int(self.config.rkv_redundancy_window) window = int(slots.shape[1]) if configured_window == 0 else min(configured_window, int(slots.shape[1])) if window > 0: - k_cache, _ = self.get_layer_kv_cache(layer_idx) window_slots = slots[:, -window:] - window_keys = k_cache.index_select(0, window_slots.reshape(-1)).view( - len(seqs), - window, - k_cache.shape[1], - k_cache.shape[2], + window_keys = self.materialize_attention_keys( + layer_idx, + window_slots, ) redundancy = self.redundancy_scores_from_keys_batch( window_keys, diff --git a/src/sparsevllm/engine/cache_manager/snapkv.py b/src/sparsevllm/engine/cache_manager/snapkv.py index baa5e759..ac4fe079 100644 --- a/src/sparsevllm/engine/cache_manager/snapkv.py +++ b/src/sparsevllm/engine/cache_manager/snapkv.py @@ -21,6 +21,7 @@ from sparsevllm.utils.profiler import profiler from .base import ( + AttentionCacheWrite, CacheManager, ExplicitKVPayload, LayerBatchStates, @@ -28,6 +29,7 @@ SparseSelection, ) from .raw_kv_offload import RawKVOffloadBuffer +from .storage import ExplicitKVStorage, create_attention_cache_storage _INT32_BYTES = 4 @@ -93,6 +95,15 @@ def resolve_snapkv_cache_capacity( class SnapKVCacheManager(CacheManager): def __init__(self, config: Config, parallel_context: ParallelContext): super().__init__(config, parallel_context) + self.attention_cache_storage = ( + create_attention_cache_storage( + config, + num_kv_heads=self.num_kv_heads, + head_dim=self.head_dim, + ) + if config.pyramid_layer_ratios is None + else None + ) self.pyramidkv_prefill_staging_num_slots = 0 self.pyramidkv_prefill_staging_kv_cache = None self._pyramidkv_prefill_staging_active = False @@ -381,15 +392,27 @@ def allocate_kv_cache(self): f"{config.num_kvcache_slots} tokens, " f"row_slot_map_bytes={row_slot_map_bytes}." ) - self.kv_cache = torch.empty( - 2, - num_layers, - config.num_kvcache_slots, - self.num_kv_heads, - self.head_dim, - dtype=self.hf_config.torch_dtype, + storage = self.attention_cache_storage + if storage is None: + raise RuntimeError( + "Uniform SnapKV requires an attention cache storage." + ) + storage.allocate( + num_layers=num_layers, + num_slots=config.num_kvcache_slots, device=self.device, ) + self.kv_cache = ( + storage.cache + if isinstance(storage, ExplicitKVStorage) + else None + ) + + def attention_cache_bytes_per_slot_per_layer(self) -> int: + storage = getattr(self, "attention_cache_storage", None) + if storage is None: + return super().attention_cache_bytes_per_slot_per_layer() + return int(storage.bytes_per_slot_per_layer()) def get_layer_batch_states(self, layer_idx: int) -> LayerBatchStates: self.kv_layer_index(layer_idx) @@ -404,6 +427,79 @@ def get_layer_kv_cache(self, layer_idx: int) -> tuple[torch.Tensor, torch.Tensor else: raise ValueError + def store_attention_payload( + self, + layer_idx: int, + payload: AttentionCacheWrite, + ) -> torch.Tensor: + storage = getattr(self, "attention_cache_storage", None) + if storage is None or isinstance(storage, ExplicitKVStorage): + return super().store_attention_payload(layer_idx, payload) + slot_mapping = self.layer_batch_states[layer_idx].slot_mapping + if slot_mapping is None: + raise RuntimeError( + f"Attention cache store requires slot_mapping at layer={layer_idx}." + ) + storage.store( + self.kv_layer_index(layer_idx), + slot_mapping, + payload, + ) + return slot_mapping + + def get_layer_compute_payload( + self, + layer_idx: int, + active_slots: torch.Tensor, + req_indices: torch.Tensor, + context_lens: torch.Tensor, + selection: SparseSelection | None = None, + ): + storage = getattr(self, "attention_cache_storage", None) + if storage is None or isinstance(storage, ExplicitKVStorage): + return super().get_layer_compute_payload( + layer_idx, + active_slots, + req_indices, + context_lens, + selection, + ) + return ( + storage.layer_payload(self.kv_layer_index(layer_idx)), + active_slots, + req_indices, + context_lens, + ) + + def get_prefill_compute_payload( + self, + layer_idx: int, + k_current: torch.Tensor, + v_current: torch.Tensor, + selection: SparseSelection, + active_slots: torch.Tensor, + req_indices: torch.Tensor, + context_lens: torch.Tensor, + ): + storage = getattr(self, "attention_cache_storage", None) + if storage is None or isinstance(storage, ExplicitKVStorage): + return super().get_prefill_compute_payload( + layer_idx, + k_current, + v_current, + selection, + active_slots, + req_indices, + context_lens, + ) + return self.get_layer_compute_payload( + layer_idx, + active_slots, + req_indices, + context_lens, + selection, + ) + def get_layer_store_view(self, layer_idx: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: if self.has_prefill_staging_view(layer_idx): return ( @@ -2519,6 +2615,7 @@ def _prepare_decode_static_uniform( state.max_context_len = real_max_context_len state.req_indices = req_indices self._decode_static_state_binding_key = None + self.validate_decode_cuda_graph_slot_mappings() return input_ids, positions, None def _allocate_decode_batch_all_layers( @@ -2707,6 +2804,7 @@ def prepare_decode_static( layers_req_indices, max_context_lens, ) + self.validate_decode_cuda_graph_slot_mappings() first_layer = int(self.kv_transformer_layer_indices()[0]) slot_mapping.copy_(layers_slot_mapping[first_layer]) diff --git a/src/sparsevllm/engine/cache_manager/standard.py b/src/sparsevllm/engine/cache_manager/standard.py index d603e34a..74a26780 100644 --- a/src/sparsevllm/engine/cache_manager/standard.py +++ b/src/sparsevllm/engine/cache_manager/standard.py @@ -1505,6 +1505,6 @@ def prepare_decode_static( self.layer_batch_state.context_lens = context_lens self.layer_batch_state.max_context_len = int(real_context_lens.max()) if real_batch_size > 0 else 0 self.layer_batch_state.req_indices = req_indices - self._validate_attention_slot_mapping(slot_mapping) + self.validate_decode_cuda_graph_slot_mappings() return input_ids, positions, None diff --git a/src/sparsevllm/engine/cache_manager/storage/base.py b/src/sparsevllm/engine/cache_manager/storage/base.py index f64726ef..784d9e37 100644 --- a/src/sparsevllm/engine/cache_manager/storage/base.py +++ b/src/sparsevllm/engine/cache_manager/storage/base.py @@ -31,6 +31,11 @@ def layer_payload(self, layer_idx: int) -> AttentionPayload: ... def validate_slot_mapping(self, slot_mapping: torch.Tensor) -> None: ... + def validate_slot_mappings( + self, + slot_mappings: tuple[torch.Tensor, ...], + ) -> None: ... + def store( self, layer_idx: int, @@ -38,6 +43,15 @@ def store( payload: AttentionCacheWrite, ) -> None: ... + def copy_slots( + self, + layer_idx: int, + source_slots: torch.Tensor, + destination_slots: torch.Tensor, + ) -> None: ... + + def slot_capacity(self) -> int: ... + def bytes_per_slot_per_layer(self) -> int: ... def accounting_tensors(self) -> tuple[torch.Tensor, ...]: ... diff --git a/src/sparsevllm/engine/cache_manager/storage/explicit_kv.py b/src/sparsevllm/engine/cache_manager/storage/explicit_kv.py index f2db9734..596c1b8e 100644 --- a/src/sparsevllm/engine/cache_manager/storage/explicit_kv.py +++ b/src/sparsevllm/engine/cache_manager/storage/explicit_kv.py @@ -91,6 +91,13 @@ def validate_slot_mapping(self, slot_mapping: torch.Tensor) -> None: f"slots={slot_mapping.device} cache={cache.device}." ) + def validate_slot_mappings( + self, + slot_mappings: tuple[torch.Tensor, ...], + ) -> None: + for slot_mapping in slot_mappings: + self.validate_slot_mapping(slot_mapping) + def store( self, layer_idx: int, @@ -149,5 +156,50 @@ def bytes_per_slot_per_layer(self) -> int: element_size = torch.tensor([], dtype=self.dtype).element_size() return int(2 * self.num_kv_heads * self.head_dim * element_size) + def slot_capacity(self) -> int: + return int(self._require_cache().shape[2]) + + @torch.no_grad() + def copy_slots( + self, + layer_idx: int, + source_slots: torch.Tensor, + destination_slots: torch.Tensor, + ) -> None: + payload = self.layer_payload(layer_idx) + source_slots = source_slots.to( + device=payload.k_cache.device, + dtype=torch.long, + ).reshape(-1) + destination_slots = destination_slots.to( + device=payload.k_cache.device, + dtype=torch.long, + ).reshape(-1) + if source_slots.shape != destination_slots.shape: + raise ValueError( + "Explicit KV slot copy requires equal source/destination shapes, " + f"got {tuple(source_slots.shape)} and " + f"{tuple(destination_slots.shape)}." + ) + if source_slots.numel() == 0: + return + slot_count = int(payload.k_cache.shape[0]) + in_bounds = ( + (source_slots >= 0) + & (source_slots < slot_count) + & (destination_slots >= 0) + & (destination_slots < slot_count) + ).all() + if in_bounds.is_cuda: + torch._assert_async(in_bounds) + elif not bool(in_bounds.item()): + raise ValueError( + f"Explicit KV slot copy indices must be in [0, {slot_count})." + ) + k_selected = payload.k_cache.index_select(0, source_slots) + v_selected = payload.v_cache.index_select(0, source_slots) + payload.k_cache.index_copy_(0, destination_slots, k_selected) + payload.v_cache.index_copy_(0, destination_slots, v_selected) + def accounting_tensors(self) -> tuple[torch.Tensor, ...]: return (self._require_cache(),) diff --git a/src/sparsevllm/engine/cache_manager/storage/mla_latent.py b/src/sparsevllm/engine/cache_manager/storage/mla_latent.py index 3de9575d..048efd45 100644 --- a/src/sparsevllm/engine/cache_manager/storage/mla_latent.py +++ b/src/sparsevllm/engine/cache_manager/storage/mla_latent.py @@ -39,8 +39,9 @@ def __init__( ) self.latent_cache: torch.Tensor | None = None self.rope_cache: torch.Tensor | None = None - self._validated_slot_mapping_key: tuple[str, int, int] | None = None - self._validated_store_calls_remaining = 0 + self._validated_store_calls_remaining: dict[ + tuple[str, int, int], int + ] = {} def allocate( self, @@ -72,8 +73,7 @@ def allocate( dtype=self.dtype, device=device, ) - self._validated_slot_mapping_key = None - self._validated_store_calls_remaining = 0 + self._validated_store_calls_remaining.clear() def _require_caches(self) -> tuple[torch.Tensor, torch.Tensor]: if self.latent_cache is None or self.rope_cache is None: @@ -100,7 +100,7 @@ def _slot_mapping_key(slot_mapping: torch.Tensor) -> tuple[str, int, int]: int(slot_mapping.numel()), ) - def validate_slot_mapping(self, slot_mapping: torch.Tensor) -> None: + def _validate_slot_mapping(self, slot_mapping: torch.Tensor) -> None: latent_cache, _ = self._require_caches() if slot_mapping.ndim != 1: raise ValueError( @@ -119,8 +119,28 @@ def validate_slot_mapping(self, slot_mapping: torch.Tensor) -> None: slot_mapping, cache_slot_count=int(latent_cache.shape[1]), ) - self._validated_slot_mapping_key = self._slot_mapping_key(slot_mapping) - self._validated_store_calls_remaining = int(latent_cache.shape[0]) + + def validate_slot_mappings( + self, + slot_mappings: tuple[torch.Tensor, ...], + ) -> None: + if not slot_mappings: + raise ValueError("MLA slot mapping validation requires at least one layer.") + remaining: dict[tuple[str, int, int], int] = {} + validated: set[tuple[str, int, int]] = set() + for slot_mapping in slot_mappings: + key = self._slot_mapping_key(slot_mapping) + if key not in validated: + self._validate_slot_mapping(slot_mapping) + validated.add(key) + remaining[key] = remaining.get(key, 0) + 1 + self._validated_store_calls_remaining = remaining + + def validate_slot_mapping(self, slot_mapping: torch.Tensor) -> None: + latent_cache, _ = self._require_caches() + self.validate_slot_mappings( + (slot_mapping,) * int(latent_cache.shape[0]) + ) def store( self, @@ -135,10 +155,11 @@ def store( ) destination = self.layer_payload(layer_idx) slot_mapping_key = self._slot_mapping_key(slot_mapping) - use_prevalidated_mapping = ( - slot_mapping_key == self._validated_slot_mapping_key - and self._validated_store_calls_remaining > 0 + remaining = self._validated_store_calls_remaining.get( + slot_mapping_key, + 0, ) + use_prevalidated_mapping = remaining > 0 copy_latent_to_cache( payload.latent, payload.rope, @@ -148,13 +169,70 @@ def store( validate_slots=not use_prevalidated_mapping, ) if use_prevalidated_mapping: - self._validated_store_calls_remaining -= 1 - if self._validated_store_calls_remaining == 0: - self._validated_slot_mapping_key = None + if remaining == 1: + del self._validated_store_calls_remaining[slot_mapping_key] + else: + self._validated_store_calls_remaining[slot_mapping_key] = ( + remaining - 1 + ) def bytes_per_slot_per_layer(self) -> int: element_size = torch.tensor([], dtype=self.dtype).element_size() return int((self.kv_lora_rank + self.rope_dim) * element_size) + def slot_capacity(self) -> int: + latent_cache, _ = self._require_caches() + return int(latent_cache.shape[1]) + + @torch.no_grad() + def copy_slots( + self, + layer_idx: int, + source_slots: torch.Tensor, + destination_slots: torch.Tensor, + ) -> None: + payload = self.layer_payload(layer_idx) + source_slots = source_slots.to( + device=payload.latent_cache.device, + dtype=torch.long, + ).reshape(-1) + destination_slots = destination_slots.to( + device=payload.latent_cache.device, + dtype=torch.long, + ).reshape(-1) + if source_slots.shape != destination_slots.shape: + raise ValueError( + "MLA latent slot copy requires equal source/destination shapes, " + f"got {tuple(source_slots.shape)} and " + f"{tuple(destination_slots.shape)}." + ) + if source_slots.numel() == 0: + return + slot_count = int(payload.latent_cache.shape[0]) + in_bounds = ( + (source_slots >= 0) + & (source_slots < slot_count) + & (destination_slots >= 0) + & (destination_slots < slot_count) + ).all() + if in_bounds.is_cuda: + torch._assert_async(in_bounds) + elif not bool(in_bounds.item()): + raise ValueError( + f"MLA latent slot copy indices must be in [0, {slot_count})." + ) + latent_selected = payload.latent_cache.index_select(0, source_slots) + rope_selected = payload.rope_cache.index_select(0, source_slots) + payload.latent_cache.index_copy_( + 0, + destination_slots, + latent_selected, + ) + payload.rope_cache.index_copy_( + 0, + destination_slots, + rope_selected, + ) + def accounting_tensors(self) -> tuple[torch.Tensor, ...]: return self._require_caches() diff --git a/src/sparsevllm/engine/chain_cache.py b/src/sparsevllm/engine/chain_cache.py index ae438f6f..02e07ad7 100644 --- a/src/sparsevllm/engine/chain_cache.py +++ b/src/sparsevllm/engine/chain_cache.py @@ -13,7 +13,7 @@ CHAIN_PREFIX_METHODS = frozenset( - {"snapkv", "h2o", "pyramidkv", "rkv", "skipkv"} + {"streamingllm", "snapkv", "h2o", "pyramidkv", "rkv", "skipkv"} ) RADIX_PREFIX_METHODS = frozenset({"", "omnikv", "quest"}) PREFIX_CACHE_MODES = frozenset({"auto", "radix", "chain"}) @@ -195,6 +195,10 @@ def build_chain_cache_fingerprint(config: Any) -> bytes: hf_config = getattr(config, "hf_config", None) method = str(getattr(config, "vllm_sparse_method", "") or "") method_fields = { + "streamingllm": ( + "num_sink_tokens", + "num_recent_tokens", + ), "snapkv": ( "num_sink_tokens", "num_recent_tokens", diff --git a/src/sparsevllm/engine/decode_cuda_graph.py b/src/sparsevllm/engine/decode_cuda_graph.py index 665b934f..1620262b 100644 --- a/src/sparsevllm/engine/decode_cuda_graph.py +++ b/src/sparsevllm/engine/decode_cuda_graph.py @@ -111,6 +111,10 @@ def __init__( self.last_state_key: DecodeCudaGraphKey | None = None self.last_real_batch_size: int | None = None self.graph_pool = graph_pool + self.capture_count = 0 + self.replay_count = 0 + self.eager_static_count = 0 + self.force_eager_count = 0 def _resolve_max_cached_graphs(self) -> int | None: resolver = getattr(self.cache_manager, "decode_cuda_graph_max_cached_graphs", None) @@ -421,6 +425,11 @@ def _capture( _ = logits.argmax(dim=-1) self.platform.synchronize() + # Static metadata was validated before the eager warmup. Some cache + # layouts use that validation to suppress graph-unsafe host checks for + # exactly one model forward, so establish a fresh scope for capture. + self.cache_manager.validate_decode_cuda_graph_slot_mappings() + with profiler.record("decode_cuda_graph_capture"): self.sparse_controller.prepare_forward(seqs, is_prefill=False) # Dynamic score paths can replace Python state fields during the @@ -469,6 +478,7 @@ def _capture( if sparse_keepalive is not None: keepalive.extend(sparse_keepalive()) state.keepalive = keepalive + self.capture_count = int(getattr(self, "capture_count", 0)) + 1 return state def run( @@ -493,6 +503,9 @@ def run( real_batch_size = len(seqs) force_eager = getattr(self.cache_manager, "decode_cuda_graph_force_eager", None) if force_eager is not None and force_eager(): + self.force_eager_count = int( + getattr(self, "force_eager_count", 0) + ) + 1 return self.run_eager_static(seqs), None graph_batch_size = self._select_graph_batch_size(real_batch_size) @@ -515,6 +528,7 @@ def run( self._restore_sparse_state_refs(state) with profiler.record("decode_cuda_graph_replay_after_capture"): state.graph.replay() + self.replay_count = int(getattr(self, "replay_count", 0)) + 1 logits = state.logits[:real_batch_size] if state.logits is not None else None token_ids = state.token_ids[:real_batch_size] if state.token_ids is not None else None return logits, token_ids @@ -522,6 +536,7 @@ def run( self._restore_sparse_state_refs(state) with profiler.record("decode_cuda_graph_replay"): state.graph.replay() + self.replay_count = int(getattr(self, "replay_count", 0)) + 1 logits = state.logits[:real_batch_size] if state.logits is not None else None token_ids = state.token_ids[:real_batch_size] if state.token_ids is not None else None return logits, token_ids @@ -530,6 +545,9 @@ def run_eager_static(self, seqs: list[Sequence]) -> torch.Tensor | None: """Run decode eagerly through the same static-compatible path used by graphs.""" if not seqs: raise ValueError("static decode requires a non-empty decode batch.") + self.eager_static_count = int( + getattr(self, "eager_static_count", 0) + ) + 1 real_batch_size = len(seqs) graph_batch_size = self._select_graph_batch_size(real_batch_size) diff --git a/src/sparsevllm/engine/model_runner.py b/src/sparsevllm/engine/model_runner.py index 78b606f1..036f3d4e 100644 --- a/src/sparsevllm/engine/model_runner.py +++ b/src/sparsevllm/engine/model_runner.py @@ -783,6 +783,17 @@ def prefix_cache_set_eviction_priority( ) def debug_sparse_state_summary(self) -> dict[str, object]: + def parallel_group_summary(group) -> dict[str, object] | None: + if group is None: + return None + return { + "rank": int(group.rank), + "size": int(group.size), + "ranks": [int(rank) for rank in group.ranks], + } + + parallel_context = self.parallel_context + config = getattr(self, "config", None) moe_synced = {} moe_local = {} model = getattr(getattr(self, "model", None), "model", None) @@ -819,10 +830,73 @@ def debug_sparse_state_summary(self) -> dict[str, object]: state["mixed_prefix_cache"] = ( prefix_cache_coordinator.debug_state_summary() ) + graph_runner = getattr(self, "decode_cuda_graph_runner", None) + graph_key = getattr(graph_runner, "last_state_key", None) + graph_summary = { + "enabled": bool( + getattr(config, "decode_cuda_graph", False) + ), + "capture_count": int( + getattr(graph_runner, "capture_count", 0) + ), + "replay_count": int( + getattr(graph_runner, "replay_count", 0) + ), + "eager_static_count": int( + getattr(graph_runner, "eager_static_count", 0) + ), + "force_eager_count": int( + getattr(graph_runner, "force_eager_count", 0) + ), + "cached_graph_count": len( + getattr(graph_runner, "_graphs", {}) + ), + "last_state_key": ( + { + "method": str(graph_key.method or ""), + "batch_size": int(graph_key.batch_size), + "context_capacity": int(graph_key.context_capacity), + "is_long_text": bool(graph_key.is_long_text), + "capture_sampling": bool(graph_key.capture_sampling), + } + if graph_key is not None + else None + ), + } return { "world_rank": self.parallel_context.world_rank, "ep_rank": self.parallel_context.ep_rank, + "parallel": { + "configured": { + "tensor_parallel_size": int( + getattr(config, "tensor_parallel_size", parallel_context.tp_size) + ), + "expert_parallel_size": int( + getattr(config, "expert_parallel_size", parallel_context.ep_size) + ), + "data_parallel_size": int( + getattr(config, "data_parallel_size", parallel_context.dp_size) + ), + "world_size": int( + getattr(config, "world_size", parallel_context.world_size) + ), + }, + "effective": { + "world": parallel_group_summary(parallel_context.world), + "attention": parallel_group_summary(parallel_context.attention), + "expert": parallel_group_summary(parallel_context.expert), + "moe_tensor": parallel_group_summary( + parallel_context.moe_tensor or parallel_context.tensor + ), + "data": parallel_group_summary(parallel_context.data), + }, + "attention_replicated_for_ep": bool( + parallel_context.ep_size > 1 + and parallel_context.attention_tp_size == 1 + ), + }, "state": state, + "decode_cuda_graph": graph_summary, "last_logits": ( _debug_tensor_summary(self.debug_last_logits) if hasattr(self, "debug_last_logits") @@ -833,12 +907,14 @@ def debug_sparse_state_summary(self) -> dict[str, object]: } def debug_last_logits_cpu(self) -> torch.Tensor | None: + if self.rank != 0: + return None logits = getattr(self, "debug_last_logits", None) if logits is None: raise RuntimeError( "No debug logits are available. Set SPARSEVLLM_DEBUG_RUNTIME=1 before engine startup." ) - return logits.detach().cpu() if self.rank == 0 else None + return logits.detach().cpu() def debug_hidden_states_cpu(self) -> dict[int, torch.Tensor] | None: model = getattr(getattr(self, "model", None), "model", None) @@ -861,6 +937,8 @@ def debug_moe_states_cpu(self) -> dict[int, dict[str, torch.Tensor]] | None: snapshots = {} for layer_idx, layer in enumerate(layers): block = getattr(layer, "mlp", None) + if block is None or not hasattr(block, "experts"): + continue required = { "input": getattr(block, "debug_last_input", None), "topk_ids": getattr(block, "debug_last_topk_ids", None), @@ -922,18 +1000,29 @@ def _debug_any_mismatch_from_world_rank_zero(self, tensor: torch.Tensor) -> bool def debug_replica_consistency(self) -> dict[str, object] | None: logits = getattr(self, "debug_last_logits", None) - if logits is None: - return None - logits_max_abs, logits_tolerance_ratio = self._debug_float_error_from_world_rank_zero( - logits, - atol=0.05, - rtol=0.05, - ) - result: dict[str, object] = { - "last_logits_max_abs": logits_max_abs, - "last_logits_tolerance_ratio": logits_tolerance_ratio, - "moe_layers": {}, - } + if self.parallel_context.attention_tp_size > 1: + result: dict[str, object] = { + "last_logits_max_abs": None, + "last_logits_tolerance_ratio": None, + "last_logits_comparison": "not_applicable_tp_vocab_sharded", + "moe_layers": {}, + } + else: + if logits is None: + return None + logits_max_abs, logits_tolerance_ratio = ( + self._debug_float_error_from_world_rank_zero( + logits, + atol=0.05, + rtol=0.05, + ) + ) + result = { + "last_logits_max_abs": logits_max_abs, + "last_logits_tolerance_ratio": logits_tolerance_ratio, + "last_logits_comparison": "compared", + "moe_layers": {}, + } model = getattr(getattr(self, "model", None), "model", None) layers = getattr(model, "layers", ()) for layer_idx in sorted({0, len(layers) - 1} if layers else set()): @@ -1186,10 +1275,19 @@ def run_model(self, input_ids: torch.Tensor, positions: torch.Tensor, is_prefill _stage = 'prefill' if is_prefill else 'decode' with profiler.record(f"model_run_model_{_stage}"): logits = self.model.compute_logits(self.model(input_ids, positions)) - if os.getenv("SPARSEVLLM_DEBUG_RUNTIME", "0") == "1": - self.debug_last_logits = logits.detach().clone() + self._record_debug_logits(logits) return logits + def _record_debug_logits(self, logits: torch.Tensor | None) -> None: + if ( + os.getenv("SPARSEVLLM_DEBUG_RUNTIME", "0") == "1" + and isinstance(logits, torch.Tensor) + ): + # A clone captured inside run_model keeps the capture-time value on + # CUDA Graph replay. Refresh outside replay so debug/validation + # observes the business step that actually completed. + self.debug_last_logits = logits.detach().clone() + def run_logits_for_compare(self, seqs: list[Sequence], is_prefill: bool) -> torch.Tensor | None: """Debug logits-alignment path: execute one step and return rank-0 logits without sampling.""" try: @@ -1268,6 +1366,7 @@ def run( else: logits = self.decode_cuda_graph_runner.run_eager_static(seqs) graph_token_ids = None + self._record_debug_logits(logits) if self.rank != 0: self._post_sparse_forward(seqs, is_prefill) return None, None diff --git a/src/sparsevllm/engine/sparse_controller.py b/src/sparsevllm/engine/sparse_controller.py index 417be53f..0a349bff 100644 --- a/src/sparsevllm/engine/sparse_controller.py +++ b/src/sparsevllm/engine/sparse_controller.py @@ -3,6 +3,7 @@ import torch import torch.nn.functional as F from sparsevllm.config import Config +from sparsevllm.configs.model import resolve_attention_qk_head_dim from sparsevllm.engine.activation_controller import ActivationController from sparsevllm.engine.sequence import Sequence from sparsevllm.engine.cache_manager import CacheManager, SparseSelection @@ -83,10 +84,7 @@ def __init__(self, config: Config, cache_manager: CacheManager): self.num_sink = self.config.num_sink_tokens self.num_recent = self.config.num_recent_tokens self.decode_keep_tokens = self.config.decode_keep_tokens - head_dim = int( - getattr(self.config.hf_config, "head_dim", None) - or (self.config.hf_config.hidden_size // self.config.hf_config.num_attention_heads) - ) + head_dim = resolve_attention_qk_head_dim(self.config.hf_config) self.attn_softmax_scale = float(head_dim) ** -0.5 score_dtype_name = str(getattr(self.config, "sparse_attn_score_dtype", "float32") or "float32").lower() self.attn_score_dtype = { @@ -104,6 +102,15 @@ def __init__(self, config: Config, cache_manager: CacheManager): self.layer_batch_sparse_states[i] = LayerBatchSparseState() self._decode_attn_score_buffers: dict[int, torch.Tensor] = {} self._omnikv_decode_attn_score_buffer: torch.Tensor | None = None + self._omnikv_decode_selection_buffers: dict[ + tuple[int, tuple[int, ...], int, int, str], + tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + ], + ] = {} self._snapkv_decode_reduced_attn_score_buffers: dict[int, torch.Tensor] = {} self._h2o_decode_attn_score_buffers: dict[tuple[int, ...], torch.Tensor] = {} @@ -153,10 +160,50 @@ def decode_cuda_graph_keepalive_tensors(self) -> list[torch.Tensor]: tensors = self.activation_controller.decode_cuda_graph_keepalive_tensors() if self._omnikv_decode_attn_score_buffer is not None: tensors.append(self._omnikv_decode_attn_score_buffer) + for buffers in self._omnikv_decode_selection_buffers.values(): + tensors.extend(buffers) tensors.extend(self._snapkv_decode_reduced_attn_score_buffers.values()) tensors.extend(self._h2o_decode_attn_score_buffers.values()) return tensors + def _get_omnikv_decode_selection_buffers( + self, + *, + obs_layer_idx: int, + target_layers: list[int], + batch_size: int, + max_context_len: int, + device: torch.device, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + cache = getattr(self, "_omnikv_decode_selection_buffers", None) + if cache is None: + cache = {} + self._omnikv_decode_selection_buffers = cache + key = ( + int(obs_layer_idx), + tuple(int(layer_idx) for layer_idx in target_layers), + int(batch_size), + int(max_context_len), + str(device), + ) + buffers = cache.get(key) + if buffers is not None: + return buffers + if device.type == "cuda" and torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "OmniKV decode CUDA graph capture requires selection buffers " + "to be allocated by the eager capture warmup." + ) + selection_shape = (int(batch_size), int(max_context_len)) + buffers = ( + torch.empty(selection_shape, dtype=torch.int32, device=device), + torch.empty(selection_shape, dtype=torch.int32, device=device), + torch.empty((int(batch_size),), dtype=torch.int32, device=device), + torch.arange(int(batch_size), dtype=torch.int32, device=device), + ) + cache[key] = buffers + return buffers + def _debug_record_dynamic_selection(self, bucket: str, layer_idx: int, **fields): entry = self.debug_dynamic_selection.setdefault(bucket, {}).setdefault(str(int(layer_idx)), {"calls": 0}) entry["calls"] += 1 @@ -1444,6 +1491,13 @@ def _streamingllm_prefill_eviction(self, seqs: list[Sequence]): by_kv_len.setdefault(int(kv_len), []).append((b_idx, seq)) for kv_len, group in by_kv_len.items(): + if log_level == 'DEBUG': + for _b_idx, seq in group: + logger.debug( + "[StreamingLLM] prefill eviction: " + f"layer={layer_idx} seq_id={seq.seq_id} " + f"kv_len={kv_len} budget={budget}" + ) group_seqs = [seq for _b_idx, seq in group] if free_prefix_recent is not None: key = (tuple(int(seq.seq_id) for seq in group_seqs), int(kv_len)) @@ -1793,7 +1847,6 @@ def _update_dynamic_omnikv_indices(self, obs_layer_idx, target_layers): # 4. 根据方法更新目标层状态 if self.sparse_method == 'omnikv': - local_req_indices = torch.arange(batch_size, dtype=torch.int32, device=self.device) decode_keep = int(decode_keep) k_max = min(decode_keep, int(search_scores.size(1))) if k_max > 0: @@ -1810,6 +1863,46 @@ def _update_dynamic_omnikv_indices(self, obs_layer_idx, target_layers): max_recent_or_chunk = int(self.num_recent) max_sparse_context_len = int(self.num_sink) + int(k_max) + max_recent_or_chunk slot_source_layer = int(target_layers[0]) + graph_outputs = None + if ( + not ctx.is_prefill + and bool( + getattr( + getattr(self, "config", None), + "decode_cuda_graph", + False, + ) + ) + ): + graph_outputs = self._get_omnikv_decode_selection_buffers( + obs_layer_idx=obs_layer_idx, + target_layers=target_layers, + batch_size=batch_size, + max_context_len=max_sparse_context_len, + device=token_scores.device, + ) + ( + keep_indices_out, + active_slots_out, + new_context_lens_out, + local_req_indices, + ) = graph_outputs + else: + keep_indices_out = None + active_slots_out = None + new_context_lens_out = None + local_req_indices = torch.arange( + batch_size, + dtype=torch.int32, + device=self.device, + ) + output_kwargs = {} + if graph_outputs is not None: + output_kwargs = { + "keep_indices_out": keep_indices_out, + "active_slots_out": active_slots_out, + "new_context_lens_out": new_context_lens_out, + } keep_indices, active_slots, new_context_lens = build_omnikv_keep_and_slots( topk_indices, topk_lens, @@ -1819,7 +1912,17 @@ def _update_dynamic_omnikv_indices(self, obs_layer_idx, target_layers): obs_sparse_state.req_indices, self.num_sink, max_s=max_sparse_context_len, + **output_kwargs, ) + if graph_outputs is not None and ( + keep_indices is not keep_indices_out + or active_slots is not active_slots_out + or new_context_lens is not new_context_lens_out + ): + raise RuntimeError( + "OmniKV decode CUDA graph selection builder did not " + "preserve caller-owned output buffers." + ) for l_idx in target_layers: target_sparse_state = self.layer_batch_sparse_states[l_idx] diff --git a/src/sparsevllm/layers/mla_attention.py b/src/sparsevllm/layers/mla_attention.py index 53f37765..1d92826e 100644 --- a/src/sparsevllm/layers/mla_attention.py +++ b/src/sparsevllm/layers/mla_attention.py @@ -1,10 +1,12 @@ from __future__ import annotations from dataclasses import dataclass +from typing import Callable import torch from sparsevllm.engine.cache_manager.base import ( + AttentionKeyComputeView, AttentionViewMeta, DecodeComputeView, ExplicitKVPayload, @@ -329,7 +331,7 @@ def device(self) -> torch.device: def _require_mla_payload( self, - view: PrefillComputeView | DecodeComputeView, + view: PrefillComputeView | DecodeComputeView | AttentionKeyComputeView, *, operation: str, ) -> MlaLatentPayload: @@ -613,6 +615,82 @@ def bind_prefill_kv( expanded_v=expanded_v, ) + @torch.no_grad() + def materialize_expanded_keys( + self, + view: AttentionKeyComputeView, + *, + project_latent: Callable[[torch.Tensor], torch.Tensor], + ) -> torch.Tensor: + """Reconstruct the exact post-RoPE per-head keys for selected slots.""" + + if not isinstance(view, AttentionKeyComputeView): + raise TypeError( + "MLA key materialization requires AttentionKeyComputeView, got " + f"{type(view).__name__}." + ) + payload = self._require_mla_payload( + view, + operation="MLA key materialization", + ) + slots = view.active_slots + if slots.ndim == 0: + raise ValueError("MLA key materialization slots must not be scalar.") + if slots.dtype not in (torch.int32, torch.int64): + raise TypeError( + "MLA key materialization slots must use int32 or int64, got " + f"{slots.dtype}." + ) + if slots.device != self.device: + raise ValueError( + "MLA key materialization slots are on the wrong device: " + f"slots={slots.device} expected={self.device}." + ) + + flat_slots = slots.to(torch.long).reshape(-1) + latent = payload.latent_cache.index_select(0, flat_slots).squeeze(1) + rope = payload.rope_cache.index_select(0, flat_slots).squeeze(1) + projected = project_latent(latent) + qk_nope_head_dim = int(self.spec.qk_head_dim) - int(self.spec.rope_dim) + projected_width = qk_nope_head_dim + int(self.spec.value_head_dim) + expected_shape = ( + int(flat_slots.numel()), + int(self.spec.local_q_heads) * projected_width, + ) + if tuple(projected.shape) != expected_shape: + raise RuntimeError( + "MLA key projection returned an invalid shape: " + f"got={tuple(projected.shape)} expected={expected_shape}." + ) + if projected.device != self.device: + raise RuntimeError( + "MLA key projection returned the wrong device: " + f"got={projected.device} expected={self.device}." + ) + if projected.dtype != self.spec.activation_dtype: + raise TypeError( + "MLA key projection returned the wrong dtype: " + f"got={projected.dtype} expected={self.spec.activation_dtype}." + ) + + expanded = projected.view( + int(flat_slots.numel()), + self.spec.local_q_heads, + projected_width, + ) + expanded_k_nope = expanded[..., :qk_nope_head_dim] + expanded_rope = rope[:, None, :].expand( + -1, + self.spec.local_q_heads, + -1, + ) + keys = torch.cat((expanded_k_nope, expanded_rope), dim=-1) + return keys.view( + *slots.shape, + self.spec.local_q_heads, + self.spec.qk_head_dim, + ) + def run_prefill( self, q: torch.Tensor, @@ -694,7 +772,28 @@ def run_prefill( query_tokens=query_tokens, ) - explicit_view = PrefillComputeView( + explicit_view = self.build_prefill_explicit_view(workset) + return self.prefill_backend.run_prefill( + q, + explicit_view, + b_start_loc=b_start_loc, + chunk_lens=chunk_lens, + max_input_len=history.max_context_len, + ) + + def build_prefill_explicit_view( + self, + workset: MlaPrefillWorkset, + ) -> PrefillComputeView: + """Expose the exact expanded KV view used by prefill score kernels.""" + + if not isinstance(workset, MlaPrefillWorkset): + raise TypeError( + "build_prefill_explicit_view requires MlaPrefillWorkset, got " + f"{type(workset).__name__}." + ) + history = workset.history + return PrefillComputeView( meta=AttentionViewMeta( active_slots=history.packed_slots, req_indices=history.local_req_indices, @@ -706,13 +805,6 @@ def run_prefill( v_cache=workset.expanded_v, ), ) - return self.prefill_backend.run_prefill( - q, - explicit_view, - b_start_loc=b_start_loc, - chunk_lens=chunk_lens, - max_input_len=history.max_context_len, - ) def run_decode( self, diff --git a/src/sparsevllm/method_registry.py b/src/sparsevllm/method_registry.py index 98c0d693..ecfecb0f 100644 --- a/src/sparsevllm/method_registry.py +++ b/src/sparsevllm/method_registry.py @@ -50,6 +50,7 @@ PREFIX_CACHE_SUPPORTED_METHODS = { "", + "streamingllm", "omnikv", "quest", "snapkv", @@ -89,13 +90,6 @@ class ModelRuntimeCompatibility: decode_cuda_graph_methods=frozenset(CANONICAL_SPARSE_METHODS), ) -GLM4_MOE_LITE_COMPATIBILITY = ModelRuntimeCompatibility( - sparse_methods=frozenset({""}), - prefix_cache_methods=frozenset(), - requires_eager=True, -) - - QWEN3_MOE_EP_COMPATIBILITY = ModelRuntimeCompatibility( sparse_methods=_MOE_SPARSE_METHODS, prefix_cache_methods=frozenset( @@ -137,6 +131,19 @@ class ModelRuntimeCompatibility: decode_cuda_graph_methods=MINIMAX_M2_EP_COMPATIBILITY.decode_cuda_graph_methods, ) +GLM4_MOE_LITE_EP_COMPATIBILITY = ModelRuntimeCompatibility( + sparse_methods=frozenset( + {"", "streamingllm", "snapkv", "h2o", "omnikv", "rkv"} + ), + prefix_cache_methods=frozenset( + {"", "streamingllm", "snapkv", "h2o", "omnikv", "rkv"} + ), + requires_eager=True, + decode_cuda_graph_methods=frozenset( + {"", "streamingllm", "snapkv", "h2o", "omnikv", "rkv"} + ), +) + MODEL_RUNTIME_COMPATIBILITY = { **{ (model_type, ParallelMode.STANDARD): DENSE_MODEL_COMPATIBILITY @@ -148,7 +155,8 @@ class ModelRuntimeCompatibility: ("qwen3_5_moe", ParallelMode.OUTER_TP_MOE): QWEN35_MOE_COMPATIBILITY, ("minimax_m2", ParallelMode.STANDARD): MINIMAX_M2_EP_COMPATIBILITY, ("minimax_m2", ParallelMode.OUTER_TP_MOE): MINIMAX_M2_TP_EP_COMPATIBILITY, - ("glm4_moe_lite", ParallelMode.STANDARD): GLM4_MOE_LITE_COMPATIBILITY, + ("glm4_moe_lite", ParallelMode.STANDARD): GLM4_MOE_LITE_EP_COMPATIBILITY, + ("glm4_moe_lite", ParallelMode.OUTER_TP_MOE): GLM4_MOE_LITE_EP_COMPATIBILITY, } # All shipped cache managers now expose a graph-stable decode preparation path. diff --git a/src/sparsevllm/models/glm4_moe_lite.py b/src/sparsevllm/models/glm4_moe_lite.py index 1452af1e..966de03b 100644 --- a/src/sparsevllm/models/glm4_moe_lite.py +++ b/src/sparsevllm/models/glm4_moe_lite.py @@ -9,7 +9,11 @@ from transformers import Glm4MoeLiteConfig from sparsevllm.distributed import get_parallel_context -from sparsevllm.engine.cache_manager import MlaLatentWrite +from sparsevllm.configs.model import resolve_attention_qk_head_dim +from sparsevllm.engine.cache_manager import ( + AttentionKeyComputeView, + MlaLatentWrite, +) from sparsevllm.layers.embed_head import ParallelLMHead, VocabParallelEmbedding from sparsevllm.layers.layernorm import RMSNorm from sparsevllm.layers.linear import ( @@ -58,7 +62,7 @@ def build_glm4_moe_lite_mla_attention( num_q_heads=int(config.num_attention_heads), kv_lora_rank=int(config.kv_lora_rank), rope_dim=int(config.qk_rope_head_dim), - qk_head_dim=int(config.qk_nope_head_dim + config.qk_rope_head_dim), + qk_head_dim=resolve_attention_qk_head_dim(config), value_head_dim=int(config.v_head_dim), activation_dtype=activation_dtype, cache_dtype=activation_dtype, @@ -94,7 +98,7 @@ def __init__( self.kv_lora_rank = int(config.kv_lora_rank) self.qk_nope_head_dim = int(config.qk_nope_head_dim) self.qk_rope_head_dim = int(config.qk_rope_head_dim) - self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + self.qk_head_dim = resolve_attention_qk_head_dim(config) self.v_head_dim = int(config.v_head_dim) self.proj_chunk_size = int(projection_chunk_size) if self.proj_chunk_size <= 0: @@ -146,6 +150,7 @@ def __init__( int(config.hidden_size), bias=bool(config.attention_bias), ) + self._attention_key_materializer_binding: tuple[int, int] | None = None def _project_kv_history(self, latent: torch.Tensor) -> torch.Tensor: if int(latent.shape[0]) <= self.proj_chunk_size: @@ -195,6 +200,25 @@ def _reconstruct_decode_values(self, latent_output: torch.Tensor) -> torch.Tenso v_weight.transpose(1, 2), ).transpose(0, 1) + def _materialize_attention_keys( + self, + view: AttentionKeyComputeView, + ) -> torch.Tensor: + return self.mla_attention.materialize_expanded_keys( + view, + project_latent=self._project_kv_history, + ) + + def _ensure_attention_key_materializer(self, cache_manager, layer_idx: int) -> None: + binding = (id(cache_manager), int(layer_idx)) + if self._attention_key_materializer_binding == binding: + return + cache_manager.register_attention_key_materializer( + int(layer_idx), + self._materialize_attention_keys, + ) + self._attention_key_materializer_binding = binding + def _run_attention( self, q: torch.Tensor, @@ -207,6 +231,7 @@ def _run_attention( cache_manager = context.cache_manager sparse_controller = context.sparse_controller layer_idx = int(context.now_layer_idx) + self._ensure_attention_key_materializer(cache_manager, layer_idx) slot_mapping = cache_manager.store_attention_payload( layer_idx, MlaLatentWrite( @@ -264,7 +289,7 @@ def _run_attention( cache_manager.collect_prefill_attention_score( layer_idx, q, - view, + self.mla_attention.build_prefill_explicit_view(workset), b_start_loc=b_start_loc, chunk_lens=chunk_lens, ) @@ -434,18 +459,71 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: "Glm4MoeLiteSparseMoeBlock expects [tokens, hidden], got " f"{tuple(hidden_states.shape)}." ) - if int(hidden_states.shape[0]) <= self.mlp_chunk_size: - routed = self._routed_chunk(hidden_states) + debug_enabled = os.getenv("SPARSEVLLM_DEBUG_MOE", "0") == "1" + if debug_enabled: + self.debug_last_input = hidden_states.detach().clone() + if not debug_enabled: + if int(hidden_states.shape[0]) <= self.mlp_chunk_size: + routed = self._routed_chunk(hidden_states) + else: + routed = torch.cat( + [ + self._routed_chunk(chunk) + for chunk in hidden_states.split( + self.mlp_chunk_size, + dim=0, + ) + ], + dim=0, + ) + elif int(hidden_states.shape[0]) <= self.mlp_chunk_size: + router_logits, topk_weights, topk_ids = self.gate(hidden_states) + routed = self.experts(hidden_states, topk_ids, topk_weights) else: - routed = torch.cat( - [ - self._routed_chunk(chunk) - for chunk in hidden_states.split(self.mlp_chunk_size, dim=0) - ], - dim=0, + router_logits_chunks = [] + topk_weights_chunks = [] + topk_ids_chunks = [] + routed_chunks = [] + for chunk in hidden_states.split(self.mlp_chunk_size, dim=0): + router_logits, topk_weights, topk_ids = self.gate(chunk) + routed_chunks.append( + self.experts(chunk, topk_ids, topk_weights) + ) + router_logits_chunks.append(router_logits) + topk_weights_chunks.append(topk_weights) + topk_ids_chunks.append(topk_ids) + routed = torch.cat(routed_chunks, dim=0) + router_logits = torch.cat(router_logits_chunks, dim=0) + topk_weights = torch.cat(topk_weights_chunks, dim=0) + topk_ids = torch.cat(topk_ids_chunks, dim=0) + if debug_enabled: + self.debug_last_router_logits = router_logits.detach().clone() + self.debug_last_topk_ids = topk_ids.detach().clone() + self.debug_last_topk_weights = topk_weights.detach().clone() + self.debug_last_local_output = routed.detach().clone() + local_mask = (topk_ids >= self.experts.local_expert_start) & ( + topk_ids < self.experts.local_expert_end ) + local_hit_count = local_mask.sum() + self.debug_last_local_hit_count = ( + local_hit_count + if torch.cuda.is_available() + and torch.cuda.is_current_stream_capturing() + else int(local_hit_count.item()) + ) + # Each rank owns one (expert, intermediate) shard. Summing across the + # world composes pure TP, pure EP, and outer-TP x MoE-EP identically to + # the shared PackedMoeExperts contract used by Qwen3-MoE. self.parallel_context.world_all_reduce(routed) - return routed + self.shared_experts(hidden_states) + if debug_enabled: + self.debug_last_routed_output = routed.detach().clone() + output = routed + self.shared_experts(hidden_states) + if debug_enabled: + # ModelRunner's cross-rank evidence contract consumes the final + # MoE block output, including both the synced routed experts and + # the replicated shared experts. + self.debug_last_output = output.detach().clone() + return output class Glm4MoeLiteDecoderLayer(nn.Module): @@ -459,6 +537,7 @@ def __init__( decode_cuda_graph: bool, ) -> None: super().__init__() + self.parallel_context = get_parallel_context() self.self_attn = Glm4MoeLiteAttention( config, mla_attention, @@ -507,6 +586,10 @@ def forward( else: hidden_states, residual = self.input_layernorm(hidden_states, residual) hidden_states = self.self_attn(positions, hidden_states, rotary_emb) + if self.parallel_context.tp_size == 1 and self.parallel_context.ep_size > 1: + # Replicated MLA must enter the post-attention norm identically on + # every expert rank before routed experts make their next decision. + self.parallel_context.ep_broadcast(hidden_states, src_ep_rank=0) hidden_states, residual = self.post_attention_layernorm( hidden_states, residual, diff --git a/src/sparsevllm/models/spec.py b/src/sparsevllm/models/spec.py index b9b41c72..dc4e9892 100644 --- a/src/sparsevllm/models/spec.py +++ b/src/sparsevllm/models/spec.py @@ -161,6 +161,8 @@ def validate_sharding(self, hf_config: Any, topology: ParallelTopology) -> None: ), "glm4_moe_lite": ModelSpec( "GLM-4.7-Flash", + supports_expert_parallel=True, + supports_outer_tp_moe=True, runtime_class_name="Glm4MoeLiteForCausalLM", attention_tp_fields=("num_attention_heads", "vocab_size"), num_experts_field="n_routed_experts", diff --git a/src/sparsevllm/operators/mla_attention.py b/src/sparsevllm/operators/mla_attention.py index 84ca5003..d5c6aa86 100644 --- a/src/sparsevllm/operators/mla_attention.py +++ b/src/sparsevllm/operators/mla_attention.py @@ -128,10 +128,15 @@ def __init__( self.device = self.workspace.block_size.device self._validated_decode_metadata: tuple[ object, - torch.Tensor, - torch.Tensor, - torch.Tensor, - int, + list[ + tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + int, + int | None, + ] + ], ] | None = None @classmethod @@ -154,6 +159,10 @@ def supports( return SupportResult.no("platform does not support Triton") if not caps.supports_bfloat16: return SupportResult.no("device does not support BF16") + if spec.cuda_graph and not caps.supports_graph_capture: + return SupportResult.no( + "decode CUDA Graph requires platform graph capture support" + ) if spec.activation_dtype != torch.bfloat16: return SupportResult.no( f"requires BF16 activations, got {spec.activation_dtype}" @@ -184,8 +193,6 @@ def supports( return SupportResult.no( f"requires tensor parallel size 1, 2, or 4, got {spec.tp_size}" ) - if spec.cuda_graph: - return SupportResult.no("CUDA Graph execution is not validated") return SupportResult.yes() def _validate_run_inputs( @@ -282,21 +289,27 @@ def run( ) cache_slot_count = int(payload.latent_cache.shape[0]) metadata_key = ( - validation_scope, view.meta.active_slots, view.meta.req_indices, view.meta.context_lens, cache_slot_count, + view.meta.max_context_len, + ) + cached = self._validated_decode_metadata + cached_entries = ( + cached[1] + if validation_scope is not None + and cached is not None + and cached[0] is validation_scope + else [] ) - cached_key = self._validated_decode_metadata - metadata_is_validated = ( - validation_scope is not None - and cached_key is not None - and cached_key[0] is validation_scope - and cached_key[1] is metadata_key[1] - and cached_key[2] is metadata_key[2] - and cached_key[3] is metadata_key[3] - and cached_key[4] == metadata_key[4] + metadata_is_validated = any( + entry[0] is metadata_key[0] + and entry[1] is metadata_key[1] + and entry[2] is metadata_key[2] + and entry[3] == metadata_key[3] + and entry[4] == metadata_key[4] + for entry in cached_entries ) if not metadata_is_validated: validate_mla_decode_metadata( @@ -304,10 +317,16 @@ def run( view.meta.req_indices, view.meta.context_lens, cache_slot_count=cache_slot_count, + max_context_len=view.meta.max_context_len, ) - self._validated_decode_metadata = ( - metadata_key if validation_scope is not None else None - ) + if validation_scope is None: + self._validated_decode_metadata = None + else: + cached_entries.append(metadata_key) + self._validated_decode_metadata = ( + validation_scope, + cached_entries, + ) return run_mla_decode( q_nope_absorbed, q_rope, @@ -319,6 +338,8 @@ def run( output, self.workspace, softmax_scale=self.spec.softmax_scale, + attn_score=view.meta.attn_score, + max_context_len=view.meta.max_context_len, config=self.launch_config, validate_metadata=False, ) diff --git a/src/sparsevllm/triton_kernel/mla/decode_schedule.py b/src/sparsevllm/triton_kernel/mla/decode_schedule.py index 8790e38a..c11f6b52 100644 --- a/src/sparsevllm/triton_kernel/mla/decode_schedule.py +++ b/src/sparsevllm/triton_kernel/mla/decode_schedule.py @@ -243,6 +243,7 @@ def validate_mla_decode_metadata( context_lens: torch.Tensor, *, cache_slot_count: int, + max_context_len: int | None = None, ) -> None: """Synchronously validate one decode view before per-layer reuse.""" @@ -274,6 +275,17 @@ def validate_mla_decode_metadata( ) if cache_slot_count <= 0: raise ValueError("cache_slot_count must be positive") + context_capacity = ( + int(active_slots.shape[1]) + if max_context_len is None + else int(max_context_len) + ) + if not 0 < context_capacity <= int(active_slots.shape[1]): + raise ValueError( + "MLA decode max_context_len must be within the active-slot width: " + f"max_context_len={context_capacity} " + f"active_slot_width={int(active_slots.shape[1])}." + ) request_rows = request_indices.tolist() lengths = context_lens.tolist() @@ -281,10 +293,10 @@ def validate_mla_decode_metadata( for batch_index, (request_row, length) in enumerate( zip(request_rows, lengths) ): - if length < 0 or length > active_slots.shape[1]: + if length < 0 or length > context_capacity: raise ValueError( f"context_lens[{batch_index}]={length} is outside " - f"[0, {active_slots.shape[1]}]" + f"[0, {context_capacity}]" ) if request_row < 0: if length != 0: @@ -328,6 +340,8 @@ def run_mla_decode( workspace: MlaDecodeWorkspace, *, softmax_scale: float, + attn_score: torch.Tensor | None = None, + max_context_len: int | None = None, config: MlaDecodeLaunchConfig = DEFAULT_GLM_MLA_DECODE_CONFIG, validate_metadata: bool = True, ) -> torch.Tensor: @@ -371,8 +385,15 @@ def run_mla_decode( request_indices, context_lens, cache_slot_count=latent_cache.shape[0], + max_context_len=max_context_len, ) + # The 2D score path uses atomic_max across query-head tiles. Reset inside + # the captured execution so replay can never inherit a previous step's max. + # The same reset also keeps padded/tail positions explicit for 3D scores. + if attn_score is not None: + attn_score.fill_(-1.0e20) + prepare_mla_decode_schedule(context_lens, workspace, config=config) decode_stage1( q_latent, @@ -385,6 +406,8 @@ def run_mla_decode( workspace.block_size, workspace.mid_output, workspace.mid_logsumexp, + attn_score=attn_score, + max_context_len=max_context_len, softmax_scale=softmax_scale, program_count=config.program_count, block_q_heads=config.block_q_heads, diff --git a/src/sparsevllm/triton_kernel/mla/decode_stage1.py b/src/sparsevllm/triton_kernel/mla/decode_stage1.py index 55702ffe..e4b824a9 100644 --- a/src/sparsevllm/triton_kernel/mla/decode_stage1.py +++ b/src/sparsevllm/triton_kernel/mla/decode_stage1.py @@ -30,6 +30,7 @@ def _decode_stage1_kernel( context_lens, mid_output, mid_logsumexp, + attn_score, stride_slots_row, stride_slots_token, stride_q_latent_batch, @@ -49,6 +50,9 @@ def _decode_stage1_kernel( stride_mid_dim, stride_lse_head, stride_lse_block, + stride_score_batch, + stride_score_head, + stride_score_token, block_size_ptr, cache_slot_count, program_count, @@ -61,6 +65,8 @@ def _decode_stage1_kernel( BLOCK_N: tl.constexpr, PIPELINE_STAGES: tl.constexpr, MASK_HEADS: tl.constexpr, + STORE_SCORE: tl.constexpr, + REDUCE_SCORE_HEADS: tl.constexpr, ): program_id = tl.program_id(0).to(tl.int64) output_batch_start = tl.cast(0, tl.int64) @@ -161,7 +167,7 @@ def _decode_stage1_kernel( mask=valid_cache_slot[None, :], other=0.0, ) - logits = tl.dot(query_latent, cached_latent) + raw_logits = tl.dot(query_latent, cached_latent) rope_cache_offsets = ( safe_cache_slots[None, :] * stride_rope_slot @@ -172,8 +178,37 @@ def _decode_stage1_kernel( mask=valid_cache_slot[None, :], other=0.0, ) - logits += tl.dot(query_rope, cached_rope) - logits *= softmax_scale + raw_logits += tl.dot(query_rope, cached_rope) + if STORE_SCORE: + score_mask = valid_cache_slot[None, :] + if MASK_HEADS: + score_mask &= head_mask[:, None] + if REDUCE_SCORE_HEADS: + reduced_score = tl.max( + tl.where(score_mask, raw_logits, -float("inf")), + axis=0, + ) + score_offsets = ( + batch_index * stride_score_batch + + token_indices * stride_score_token + ) + tl.atomic_max( + attn_score + score_offsets, + reduced_score, + mask=valid_cache_slot, + ) + else: + score_offsets = ( + batch_index * stride_score_batch + + query_heads[:, None] * stride_score_head + + token_indices[None, :] * stride_score_token + ) + tl.store( + attn_score + score_offsets, + raw_logits, + mask=score_mask, + ) + logits = raw_logits * softmax_scale logits = tl.where( valid_cache_slot[None, :], logits, @@ -251,6 +286,8 @@ def decode_stage1( mid_output: torch.Tensor, mid_logsumexp: torch.Tensor, *, + attn_score: torch.Tensor | None = None, + max_context_len: int | None = None, softmax_scale: float, program_count: int, block_q_heads: int, @@ -276,6 +313,8 @@ def decode_stage1( "mid_output": mid_output, "mid_logsumexp": mid_logsumexp, } + if attn_score is not None: + tensors["attn_score"] = attn_score for name, tensor in tensors.items(): _require_cuda_tensor(name, tensor) if tensor.device != q_latent.device: @@ -343,6 +382,58 @@ def decode_stage1( f"mid_output has capacity for {mid_output.shape[0]} heads, " f"but {head_count} are required" ) + if attn_score is not None: + min_width = ( + int(active_slots.shape[1]) + if max_context_len is None + else int(max_context_len) + ) + if not 0 < min_width <= int(active_slots.shape[1]): + raise ValueError( + "MLA attention-score context capacity must be within the " + f"active-slot width: capacity={min_width} " + f"active_slot_width={int(active_slots.shape[1])}." + ) + if attn_score.dim() == 2: + if attn_score.dtype != torch.float32: + raise TypeError( + "Head-reduced MLA attention scores must use torch.float32, " + f"got {attn_score.dtype}." + ) + if ( + int(attn_score.shape[0]) < batch_size + or int(attn_score.shape[1]) < min_width + ): + raise ValueError( + "Head-reduced MLA attention scores must cover " + f"[batch, context]=[{batch_size}, {min_width}], got " + f"{tuple(attn_score.shape)}." + ) + elif attn_score.dim() == 3: + if attn_score.dtype not in { + torch.float32, + torch.bfloat16, + torch.float16, + }: + raise TypeError( + "Per-head MLA attention scores must use a floating dtype, " + f"got {attn_score.dtype}." + ) + if ( + int(attn_score.shape[0]) < batch_size + or int(attn_score.shape[1]) < head_count + or int(attn_score.shape[2]) < min_width + ): + raise ValueError( + "Per-head MLA attention scores must cover " + f"[batch, heads, context]=[{batch_size}, {head_count}, " + f"{min_width}], got {tuple(attn_score.shape)}." + ) + else: + raise ValueError( + "MLA attention score must be [batch, context] or " + f"[batch, heads, context], got {tuple(attn_score.shape)}." + ) if program_count <= 0: raise ValueError("program_count must be positive") if block_q_heads <= 0 or block_q_heads & (block_q_heads - 1): @@ -356,6 +447,17 @@ def decode_stage1( head_group_count = triton.cdiv(head_count, block_q_heads) mask_heads = head_count % block_q_heads != 0 + score_arg = attn_score if attn_score is not None else mid_logsumexp + if attn_score is None: + score_strides = (0, 0, 0) + elif attn_score.dim() == 2: + score_strides = ( + attn_score.stride(0), + 0, + attn_score.stride(1), + ) + else: + score_strides = attn_score.stride() _decode_stage1_kernel[(program_count,)]( q_latent, q_rope, @@ -367,6 +469,7 @@ def decode_stage1( context_lens, mid_output, mid_logsumexp, + score_arg, *active_slots.stride(), *q_latent.stride(), *q_rope.stride(), @@ -374,6 +477,7 @@ def decode_stage1( *rope_cache.stride(), *mid_output.stride(), *mid_logsumexp.stride(), + *score_strides, block_size, cache_slot_count=latent_cache.shape[0], program_count=program_count, @@ -386,6 +490,10 @@ def decode_stage1( BLOCK_N=block_n, PIPELINE_STAGES=pipeline_stages, MASK_HEADS=mask_heads, + STORE_SCORE=attn_score is not None, + REDUCE_SCORE_HEADS=( + attn_score is not None and attn_score.dim() == 2 + ), num_warps=num_warps, num_stages=1, ) diff --git a/src/sparsevllm/triton_kernel/omnikv_fused.py b/src/sparsevllm/triton_kernel/omnikv_fused.py index 7fefd513..82499166 100644 --- a/src/sparsevllm/triton_kernel/omnikv_fused.py +++ b/src/sparsevllm/triton_kernel/omnikv_fused.py @@ -87,6 +87,10 @@ def build_omnikv_keep_and_slots( req_indices: torch.Tensor, num_sink: int, max_s: int | None = None, + *, + keep_indices_out: torch.Tensor | None = None, + active_slots_out: torch.Tensor | None = None, + new_context_lens_out: torch.Tensor | None = None, ): if topk_indices.dtype != torch.int32: topk_indices = topk_indices.to(torch.int32) @@ -105,19 +109,54 @@ def build_omnikv_keep_and_slots( assert int(topk_lens.min().item()) >= 0 assert int(topk_lens.max().item()) <= k_max if max_s is None: - new_context_lens = num_sink + topk_lens + recent_chunk_lens - max_s = int(new_context_lens.max().item()) + computed_context_lens = num_sink + topk_lens + recent_chunk_lens + max_s = int(computed_context_lens.max().item()) else: max_s = int(max_s) if max_s < 0: raise ValueError(f"max_s must be >= 0, got {max_s}.") - if max_s == 0: - new_context_lens = num_sink + topk_lens + recent_chunk_lens - else: - new_context_lens = torch.empty_like(topk_lens) - keep_indices = torch.empty((batch_size, max_s), dtype=torch.int32, device=topk_indices.device) - active_slots = torch.empty((batch_size, max_s), dtype=torch.int32, device=topk_indices.device) + def _resolve_output( + name: str, + output: torch.Tensor | None, + shape: tuple[int, ...], + ) -> torch.Tensor: + if output is None: + return torch.empty( + shape, + dtype=torch.int32, + device=topk_indices.device, + ) + if output.shape != shape: + raise ValueError( + f"{name} must have shape {shape}, got {tuple(output.shape)}." + ) + if output.dtype != torch.int32 or output.device != topk_indices.device: + raise TypeError( + f"{name} must be int32 on {topk_indices.device}, got " + f"{output.dtype}/{output.device}." + ) + return output + + keep_indices = _resolve_output( + "keep_indices_out", + keep_indices_out, + (batch_size, max_s), + ) + active_slots = _resolve_output( + "active_slots_out", + active_slots_out, + (batch_size, max_s), + ) + new_context_lens = _resolve_output( + "new_context_lens_out", + new_context_lens_out, + (batch_size,), + ) + + if max_s == 0: + new_context_lens.copy_(num_sink + topk_lens + recent_chunk_lens) + return keep_indices, active_slots, new_context_lens block = 256 grid = (batch_size, triton.cdiv(max_s, block)) diff --git a/src/sparsevllm/utils/context.py b/src/sparsevllm/utils/context.py index 79678e3c..a53d850c 100644 --- a/src/sparsevllm/utils/context.py +++ b/src/sparsevllm/utils/context.py @@ -30,6 +30,7 @@ def set_context( recurrent_state_manager=None, ): global _CONTEXT + _CONTEXT.attention_validation_scope = object() _CONTEXT.is_prefill = is_prefill _CONTEXT.is_long_text = is_long_text _CONTEXT.cu_seqlens_q = cu_seqlens_q diff --git a/tests/glm_test_helpers.py b/tests/glm_test_helpers.py new file mode 100644 index 00000000..89d24793 --- /dev/null +++ b/tests/glm_test_helpers.py @@ -0,0 +1,86 @@ +from __future__ import annotations + +import hashlib +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +from sparsevllm.config import Config +from sparsevllm.distributed import ParallelContext, ParallelGroup + + +def _glm_hf_config(**overrides): + values = { + "model_type": "glm4_moe_lite", + "architectures": ["Glm4MoeLiteForCausalLM"], + "torch_dtype": torch.bfloat16, + "max_position_embeddings": 4096, + "hidden_size": 64, + "intermediate_size": 128, + "num_hidden_layers": 47, + "num_attention_heads": 20, + "num_key_value_heads": 20, + "vocab_size": 128, + "q_lora_rank": 768, + "kv_lora_rank": 512, + "qk_nope_head_dim": 192, + "qk_rope_head_dim": 64, + "v_head_dim": 256, + "moe_intermediate_size": 64, + "n_routed_experts": 64, + "n_shared_experts": 1, + "num_experts_per_tok": 4, + "n_group": 1, + "topk_group": 1, + "topk_method": "noaux_tc", + "norm_topk_prob": True, + "routed_scaling_factor": 1.8, + "mlp_layer_types": ["dense"] + ["sparse"] * 46, + "num_nextn_predict_layers": 1, + "rope_interleave": True, + "rope_scaling": None, + "attention_bias": False, + "quantization_config": None, + } + values.update(overrides) + return SimpleNamespace(**values) + + +def _glm_config(*, hf_overrides=None, **overrides) -> Config: + kwargs = { + "model": str(Path(__file__).resolve().parents[1]), + "max_model_len": 128, + "max_num_batched_tokens": 64, + "chunk_prefill_size": 64, + "enforce_eager": True, + } + kwargs.update(overrides) + with patch( + "sparsevllm.configs.runtime.AutoConfig.from_pretrained", + return_value=_glm_hf_config(**(hf_overrides or {})), + ): + return Config(**kwargs) + + +def _single_rank_parallel_context() -> ParallelContext: + singleton = ParallelGroup(None, (0,), 0, 1) + return ParallelContext( + world=singleton, + tensor=singleton, + expert=singleton, + data=singleton, + ) + + +def _tensor_sha256(tensor: torch.Tensor) -> str: + raw = ( + tensor.detach() + .contiguous() + .cpu() + .view(torch.uint8) + .numpy() + .tobytes() + ) + return hashlib.sha256(raw).hexdigest() diff --git a/tests/test_attention_cache_storage.py b/tests/test_attention_cache_storage.py index 1099e2ad..ed10a8bf 100644 --- a/tests/test_attention_cache_storage.py +++ b/tests/test_attention_cache_storage.py @@ -1,6 +1,5 @@ from __future__ import annotations -from pathlib import Path from types import SimpleNamespace from unittest.mock import patch @@ -8,7 +7,7 @@ import pytest import torch -from sparsevllm.config import Config, RuntimeLayout +from sparsevllm.config import RuntimeLayout from sparsevllm.engine.cache_manager import ( ExplicitKVPayload, ExplicitKVWrite, @@ -17,106 +16,13 @@ MlaLatentWrite, ) from sparsevllm.engine.cache_manager.standard import StandardCacheManager +from sparsevllm.engine.cache_manager.snapkv import SnapKVCacheManager from sparsevllm.engine.cache_manager.storage import ( CacheLayout, ExplicitKVStorage, MlaLatentStorage, create_attention_cache_storage, ) - - -def _glm_hf_config(**overrides): - values = { - "model_type": "glm4_moe_lite", - "architectures": ["Glm4MoeLiteForCausalLM"], - "torch_dtype": torch.bfloat16, - "max_position_embeddings": 4096, - "hidden_size": 64, - "intermediate_size": 128, - "num_hidden_layers": 47, - "num_attention_heads": 20, - "num_key_value_heads": 20, - "vocab_size": 128, - "q_lora_rank": 768, - "kv_lora_rank": 512, - "qk_nope_head_dim": 192, - "qk_rope_head_dim": 64, - "v_head_dim": 256, - "moe_intermediate_size": 64, - "n_routed_experts": 64, - "n_shared_experts": 1, - "num_experts_per_tok": 4, - "n_group": 1, - "topk_group": 1, - "topk_method": "noaux_tc", - "norm_topk_prob": True, - "routed_scaling_factor": 1.8, - "mlp_layer_types": ["dense"] + ["sparse"] * 46, - "num_nextn_predict_layers": 1, - "rope_interleave": True, - "rope_scaling": None, - "attention_bias": False, - "quantization_config": None, - } - values.update(overrides) - return SimpleNamespace(**values) - - -def _glm_config(*, hf_overrides=None, **overrides) -> Config: - kwargs = { - "model": str(Path(__file__).resolve().parents[1]), - "max_model_len": 128, - "max_num_batched_tokens": 64, - "chunk_prefill_size": 64, - "enforce_eager": True, - } - kwargs.update(overrides) - with patch( - "sparsevllm.configs.runtime.AutoConfig.from_pretrained", - return_value=_glm_hf_config(**(hf_overrides or {})), - ): - return Config(**kwargs) - - -def test_glm_config_selects_mla_latent_layout(): - config = _glm_config() - - assert config.attention_cache_layout == CacheLayout.MLA_LATENT.value - assert config.mla_prefill_workspace_bytes == 2 * 1024**3 - - -def test_mla_prefill_workspace_budget_must_be_positive(): - with pytest.raises(ValueError, match="mla_prefill_workspace_bytes"): - _glm_config(mla_prefill_workspace_bytes=0) - - -@pytest.mark.parametrize( - ("override", "message"), - [ - ({"vllm_sparse_method": "snapkv"}, "only vanilla attention"), - ({"enable_prefix_caching": True}, "does not support prefix caching"), - ({"decode_cuda_graph": True}, "does not support decode CUDA Graph"), - ({"enforce_eager": False}, "requires enforce_eager=True"), - ], -) -def test_glm_config_rejects_unsupported_storage_combinations(override, message): - with pytest.raises(NotImplementedError, match=message): - _glm_config(**override) - - -@pytest.mark.parametrize( - ("hf_override", "message"), - [ - ({"torch_dtype": torch.float16}, "requires torch.bfloat16"), - ({"kv_lora_rank": 256}, "requires kv_lora_rank=512"), - ({"qk_rope_head_dim": 32}, "qk_rope_head_dim=64"), - ], -) -def test_glm_config_rejects_unsupported_mla_storage_contract(hf_override, message): - with pytest.raises(NotImplementedError, match=message): - _glm_config(hf_overrides=hf_override) - - def test_explicit_storage_preserves_legacy_tensor_layout_and_size(): storage = ExplicitKVStorage( num_kv_heads=2, @@ -300,6 +206,42 @@ def test_storage_store_payload_types_are_not_interchangeable(): ) +@pytest.mark.parametrize("layout", ["explicit_kv", "mla_latent"]) +def test_attention_storage_copy_slots_is_overlap_safe(layout): + if layout == "explicit_kv": + storage = ExplicitKVStorage( + num_kv_heads=1, + head_dim=4, + dtype=torch.float32, + ) + else: + storage = MlaLatentStorage( + kv_lora_rank=512, + rope_dim=64, + dtype=torch.bfloat16, + ) + storage.allocate(num_layers=1, num_slots=4, device=torch.device("cpu")) + payload = storage.layer_payload(0) + tensors = ( + (payload.k_cache, payload.v_cache) + if isinstance(payload, ExplicitKVPayload) + else (payload.latent_cache, payload.rope_cache) + ) + for tensor_idx, tensor in enumerate(tensors): + for slot in range(4): + tensor[slot].fill_(tensor_idx * 10 + slot) + + storage.copy_slots( + 0, + torch.tensor([3, 1], dtype=torch.long), + torch.tensor([1, 2], dtype=torch.long), + ) + + for tensor_idx, tensor in enumerate(tensors): + assert torch.all(tensor[1] == tensor_idx * 10 + 3) + assert torch.all(tensor[2] == tensor_idx * 10 + 1) + + def test_mla_storage_reuses_one_manager_validation_across_layers(): storage = MlaLatentStorage( kv_lora_rank=512, @@ -328,6 +270,68 @@ def test_mla_storage_reuses_one_manager_validation_across_layers(): ] +def test_mla_storage_can_revalidate_between_graph_warmup_and_capture(): + storage = MlaLatentStorage( + kv_lora_rank=512, + rope_dim=64, + dtype=torch.bfloat16, + ) + storage.allocate(num_layers=2, num_slots=2, device=torch.device("cpu")) + slots = torch.tensor([0], dtype=torch.int32) + write = MlaLatentWrite( + latent=torch.empty(1, 1, 512, dtype=torch.bfloat16), + rope=torch.empty(1, 1, 64, dtype=torch.bfloat16), + ) + + with patch( + "sparsevllm.engine.cache_manager.storage.mla_latent.copy_latent_to_cache" + ) as copy: + storage.validate_slot_mapping(slots) + storage.store(0, slots, write) + storage.store(1, slots, write) + storage.validate_slot_mapping(slots) + storage.store(0, slots, write) + storage.store(1, slots, write) + + assert [call.kwargs["validate_slots"] for call in copy.call_args_list] == [ + False, + False, + False, + False, + ] + + +def test_mla_storage_prevalidates_nonuniform_layer_mappings(): + storage = MlaLatentStorage( + kv_lora_rank=512, + rope_dim=64, + dtype=torch.bfloat16, + ) + storage.allocate(num_layers=2, num_slots=2, device=torch.device("cpu")) + layer_slots = ( + torch.tensor([0], dtype=torch.int32), + torch.tensor([1], dtype=torch.int32), + ) + write = MlaLatentWrite( + latent=torch.empty(1, 1, 512, dtype=torch.bfloat16), + rope=torch.empty(1, 1, 64, dtype=torch.bfloat16), + ) + storage.validate_slot_mappings(layer_slots) + + with patch( + "sparsevllm.engine.cache_manager.storage.mla_latent.copy_latent_to_cache" + ) as copy: + storage.store(0, layer_slots[0], write) + storage.store(1, layer_slots[1], write) + storage.store(0, layer_slots[0], write) + + assert [call.kwargs["validate_slots"] for call in copy.call_args_list] == [ + False, + False, + True, + ] + + def test_standard_manager_delegates_payload_store_and_compute_view(): storage = MlaLatentStorage( kv_lora_rank=512, @@ -366,6 +370,118 @@ def test_standard_manager_delegates_payload_store_and_compute_view(): assert actual_lens is context_lens +def test_snapkv_manager_delegates_latent_store_and_compute_view(): + storage = MlaLatentStorage( + kv_lora_rank=512, + rope_dim=64, + dtype=torch.bfloat16, + ) + storage.allocate(num_layers=2, num_slots=4, device=torch.device("cpu")) + manager = object.__new__(SnapKVCacheManager) + manager.attention_cache_storage = storage + manager.runtime_layout = RuntimeLayout.dense(2) + manager.layer_batch_states = [LayerBatchStates(), LayerBatchStates()] + manager._pyramidkv_prefill_staging_active = False + manager.layer_batch_states[1].slot_mapping = torch.tensor( + [1], dtype=torch.int32 + ) + write = MlaLatentWrite( + latent=torch.empty(1, 1, 512, dtype=torch.bfloat16), + rope=torch.empty(1, 1, 64, dtype=torch.bfloat16), + ) + + with patch.object(storage, "store") as store: + returned_slots = manager.store_attention_payload(1, write) + store.assert_called_once_with( + 1, + manager.layer_batch_states[1].slot_mapping, + write, + ) + assert returned_slots is manager.layer_batch_states[1].slot_mapping + + active_slots = torch.tensor([[0, 1]], dtype=torch.int32) + req_indices = torch.tensor([0], dtype=torch.int32) + context_lens = torch.tensor([2], dtype=torch.int32) + payload, actual_slots, actual_rows, actual_lens = ( + manager.get_layer_compute_payload( + 1, + active_slots, + req_indices, + context_lens, + ) + ) + assert isinstance(payload, MlaLatentPayload) + assert actual_slots is active_slots + assert actual_rows is req_indices + assert actual_lens is context_lens + + +def test_graph_capture_prevalidates_nonuniform_latent_layer_mappings(): + storage = MlaLatentStorage( + kv_lora_rank=512, + rope_dim=64, + dtype=torch.bfloat16, + ) + storage.allocate(num_layers=2, num_slots=4, device=torch.device("cpu")) + manager = object.__new__(SnapKVCacheManager) + manager.attention_cache_storage = storage + manager.runtime_layout = RuntimeLayout.dense(2) + manager.layer_batch_states = [ + LayerBatchStates(slot_mapping=torch.tensor([0], dtype=torch.int32)), + LayerBatchStates(slot_mapping=torch.tensor([1], dtype=torch.int32)), + ] + write = MlaLatentWrite( + latent=torch.empty(1, 1, 512, dtype=torch.bfloat16), + rope=torch.empty(1, 1, 64, dtype=torch.bfloat16), + ) + + manager.validate_decode_cuda_graph_slot_mappings() + with patch( + "sparsevllm.engine.cache_manager.storage.mla_latent.copy_latent_to_cache" + ) as copy: + manager.store_attention_payload(0, write) + manager.store_attention_payload(1, write) + + assert [call.kwargs["validate_slots"] for call in copy.call_args_list] == [ + False, + False, + ] + + +def test_snapkv_explicit_compute_view_preserves_legacy_payload(): + storage = ExplicitKVStorage( + num_kv_heads=2, + head_dim=8, + dtype=torch.float16, + ) + storage.allocate(num_layers=2, num_slots=4, device=torch.device("cpu")) + manager = object.__new__(SnapKVCacheManager) + manager.attention_cache_storage = storage + manager.runtime_layout = RuntimeLayout.dense(2) + manager.kv_cache = storage.cache + manager.layer_batch_states = [LayerBatchStates(), LayerBatchStates()] + manager._pyramidkv_prefill_staging_active = False + active_slots = torch.tensor([[0, 1]], dtype=torch.int32) + req_indices = torch.tensor([0], dtype=torch.int32) + context_lens = torch.tensor([2], dtype=torch.int32) + + payload, actual_slots, actual_rows, actual_lens = ( + manager.get_layer_compute_payload( + 1, + active_slots, + req_indices, + context_lens, + ) + ) + + assert isinstance(payload, ExplicitKVPayload) + assert payload.k_cache.data_ptr() == storage.cache[0, 1].data_ptr() + assert payload.v_cache.data_ptr() == storage.cache[1, 1].data_ptr() + assert actual_slots is active_slots + assert actual_rows is req_indices + assert actual_lens is context_lens + + @pytest.mark.parametrize( ("attention_tp_size", "local_heads"), [(1, 20), (4, 5)], diff --git a/tests/test_chain_prefix_cache.py b/tests/test_chain_prefix_cache.py index 0f653143..24cd51ec 100644 --- a/tests/test_chain_prefix_cache.py +++ b/tests/test_chain_prefix_cache.py @@ -363,6 +363,7 @@ def test_chain_apply_plan_rejects_duplicate_resident_seq_owner(): ("", "auto", "radix"), ("omnikv", "auto", "radix"), ("quest", "radix", "radix"), + ("streamingllm", "auto", "chain"), ("snapkv", "auto", "chain"), ("h2o", "auto", "chain"), ("pyramidkv", "chain", "chain"), diff --git a/tests/test_compare_decode_graph_eager_logits.py b/tests/test_compare_decode_graph_eager_logits.py new file mode 100644 index 00000000..a244a0b6 --- /dev/null +++ b/tests/test_compare_decode_graph_eager_logits.py @@ -0,0 +1,246 @@ +import hashlib +from types import SimpleNamespace + +import pytest +import torch + +from scripts.debug.compare_decode_graph_eager_logits import ( + METHOD_CHOICES, + _build_method_trigger_evidence, + _build_parser, + _compare_logits, + _save_full_logits_artifact, + _start_graph_measurement, + _validate_eager_runtime, + _validate_graph_runtime, + _validate_method_trigger, +) + + +def _trace(*, row_len=4, logical_context_len=8, h2o=None, omni=False, rkv=False): + layer = {} + if omni: + layer = { + "active_slots": {"numel": 4}, + "context_lens": {"max": 4}, + } + cache = { + "live_rows": { + "0": [{"row_len": row_len}], + } + } + if h2o is not None: + cache["h2o"] = h2o + return { + "logical_context_len": logical_context_len, + "layers": {"1": layer} if layer else {}, + "cache": cache, + "rkv_materializer_layers": [0, 1] if rkv else [], + } + + +def test_graph_measurement_preserves_warmup_graph_pool_ownership(): + warmup_graph = object() + + class Runner: + def __init__(self): + self._graphs = {"warmup": SimpleNamespace(graph=warmup_graph)} + self.capture_count = 1 + self.replay_count = 1 + self.eager_static_count = 0 + self.force_eager_count = 0 + self.clear_calls = 0 + + def clear_captured_graphs(self): + self.clear_calls += 1 + self._graphs.clear() + + runner = Runner() + llm = SimpleNamespace( + model_runner=SimpleNamespace(decode_cuda_graph_runner=runner) + ) + + baseline = _start_graph_measurement(llm) + + assert runner.clear_calls == 0 + assert runner._graphs["warmup"].graph is warmup_graph + assert baseline == { + "capture_count": 1, + "replay_count": 1, + "eager_static_count": 0, + "force_eager_count": 0, + "graph_count": 1, + } + + +@pytest.mark.parametrize( + ("method", "trace", "calls"), + [ + ("vanilla", _trace(row_len=8), {}), + ( + "streamingllm", + _trace(), + {"cache.free_prefix_recent_slots_batch_layers": 1}, + ), + ("snapkv", _trace(), {"cache.free_part_slots_batch_layers": 1}), + ( + "h2o", + _trace( + h2o={ + "counters": { + "intermediate_prefill_evictions": 0, + "final_prefill_evictions": 1, + "decode_evictions": 2, + "dropped_tokens": 3, + }, + "ring_counters": {"fast_rows": 2, "fallback_rows": 0}, + } + ), + {"cache.evict_after_decode": 2}, + ), + ( + "omnikv", + _trace(row_len=8, omni=True), + {"controller._update_dynamic_omnikv_indices": 2}, + ), + ( + "rkv", + _trace(rkv=True), + { + "cache.rkv_query_attention_scores_batch": 1, + "cache.materialize_attention_keys": 2, + "cache.free_part_slots_batch_layers": 1, + }, + ), + ], +) +def test_glm_graph_method_trigger_evidence_is_machine_checkable( + method, + trace, + calls, +): + evidence = _build_method_trigger_evidence(method, [trace], calls) + + assert evidence["triggered"] is True + _validate_method_trigger(evidence) + + +def test_sparse_method_trigger_gate_rejects_unexercised_path(): + evidence = _build_method_trigger_evidence("snapkv", [_trace(row_len=8)], {}) + + with pytest.raises(RuntimeError, match="trigger gate failed"): + _validate_method_trigger(evidence) + + +def test_graph_runtime_gate_requires_capture_replay_and_zero_fallback(): + valid = { + "config_enabled": True, + "graph_active": True, + "graph_count": 1, + "capture_count": 1, + "replay_count": 3, + "eager_static_count": 0, + "force_eager_count": 0, + "counter_delta": { + "capture_count": 0, + "replay_count": 2, + "eager_static_count": 0, + "force_eager_count": 0, + }, + "fallback": False, + } + _validate_graph_runtime(valid) + + with pytest.raises(RuntimeError, match="forced eager"): + _validate_graph_runtime( + { + **valid, + "counter_delta": { + **valid["counter_delta"], + "force_eager_count": 1, + }, + } + ) + + +def test_eager_runtime_gate_rejects_a_captured_graph(): + valid = { + "config_enabled": False, + "graph_active": False, + "graph_count": 0, + "capture_count": 0, + "replay_count": 0, + "eager_static_count": 3, + "force_eager_count": 0, + "counter_delta": { + "capture_count": 0, + "replay_count": 0, + "eager_static_count": 2, + "force_eager_count": 0, + }, + } + _validate_eager_runtime(valid) + + with pytest.raises(RuntimeError, match="retained a captured"): + _validate_eager_runtime( + {**valid, "graph_active": True, "graph_count": 1} + ) + + +def test_omnikv_graph_replay_uses_tensor_selection_evidence(): + trace = _trace(row_len=8, omni=True) + trace["use_graph"] = True + + evidence = _build_method_trigger_evidence("omnikv", [trace], {}) + + assert evidence["triggered"] is True + assert evidence["execution_mode"] == "captured_replay" + + +def test_full_logits_artifact_contains_both_complete_tensors_and_hashes(tmp_path): + eager = torch.arange(24, dtype=torch.float32).reshape(3, 8) + graph = eager.clone() + path = tmp_path / "comparison.full_logits.pt" + + metadata = _save_full_logits_artifact(path, eager=eager, graph=graph) + artifact = torch.load(path, weights_only=True) + + torch.testing.assert_close(artifact["eager"], eager) + torch.testing.assert_close(artifact["graph"], graph) + assert artifact["scope"] == "all_decode_rows_and_full_vocabulary" + assert metadata["eager"]["shape"] == [3, 8] + assert metadata["eager"]["sha256"] == metadata["graph"]["sha256"] + assert metadata["artifact_sha256"] == hashlib.sha256(path.read_bytes()).hexdigest() + + +def test_full_logits_comparison_reports_tolerance_and_all_rows(): + eager = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + graph = eager + 0.01 + + result = _compare_logits(eager, graph, atol=0.02, rtol=0.0) + + assert result["within_tolerance"] is True + assert result["shape"] == [2, 2] + assert len(result["rows"]) == 2 + + +def test_cli_exposes_h2o_and_tolerance_options(tmp_path): + assert "h2o" in METHOD_CHOICES + args = _build_parser().parse_args( + [ + "--model_path", + "/checkpoint", + "--method", + "h2o", + "--output", + str(tmp_path / "result.json"), + "--atol", + "0.1", + "--rtol", + "0.2", + ] + ) + + assert args.method == "h2o" + assert args.atol == 0.1 + assert args.rtol == 0.2 diff --git a/tests/test_compare_glm_parallel_variants.py b/tests/test_compare_glm_parallel_variants.py new file mode 100644 index 00000000..2bbbaa1a --- /dev/null +++ b/tests/test_compare_glm_parallel_variants.py @@ -0,0 +1,104 @@ +from __future__ import annotations + +import json +from pathlib import Path + +import torch + +from scripts.validation.compare_glm_parallel_variants import compare_runs, main + + +def _write_run( + directory: Path, + *, + logits: torch.Tensor, + tokens: list[int], + ep_size: int, +) -> None: + directory.mkdir() + summary = { + "status": "success", + "model_path": "/checkpoint", + "seed": 7, + "prompt_len": 4, + "prompt_sha256": "prompt", + "max_tokens": len(tokens), + "sparse_method": "vanilla", + "tensor_parallel_size": 2, + "expert_parallel_size": ep_size, + "generated_token_ids": tokens, + } + (directory / "summary.json").write_text(json.dumps(summary), encoding="utf-8") + torch.save( + {"prefill": logits, "final_decode": logits}, + directory / "logits.pt", + ) + + +def test_compare_parallel_runs_accepts_bounded_drift_and_reference_tie(tmp_path): + reference = torch.tensor([[1.0, 3.0, 3.0, 0.0]]) + candidate = torch.tensor([[1.0, 2.5, 3.0, 0.0]]) + reference_dir = tmp_path / "reference" + candidate_dir = tmp_path / "candidate" + _write_run(reference_dir, logits=reference, tokens=[1, 3], ep_size=1) + _write_run(candidate_dir, logits=candidate, tokens=[2, 3], ep_size=2) + + result = compare_runs( + reference_dir, + candidate_dir, + max_abs_limit=0.5, + mean_abs_limit=0.2, + p99_abs_limit=0.5, + allow_first_token_tie=True, + tie_atol=0.0, + ) + + assert result["status"] == "success" + assert result["generated_tokens"]["first_token_tie_accepted"] is True + + +def test_compare_parallel_runs_fails_when_drift_exceeds_limit(tmp_path): + reference = torch.tensor([[1.0, 2.0]]) + candidate = torch.tensor([[1.0, 3.0]]) + reference_dir = tmp_path / "reference" + candidate_dir = tmp_path / "candidate" + _write_run(reference_dir, logits=reference, tokens=[1], ep_size=1) + _write_run(candidate_dir, logits=candidate, tokens=[1], ep_size=2) + + result = compare_runs( + reference_dir, + candidate_dir, + max_abs_limit=0.5, + mean_abs_limit=1.0, + p99_abs_limit=1.0, + allow_first_token_tie=False, + tie_atol=0.0, + ) + + assert result["status"] == "metric_failed" + assert result["gates"]["prefill_logits"] is False + + +def test_cli_writes_invalid_input_status_before_failing(tmp_path, monkeypatch): + output = tmp_path / "comparison.json" + monkeypatch.setattr( + "sys.argv", + [ + "compare_glm_parallel_variants.py", + "--reference_dir", + str(tmp_path / "missing-reference"), + "--candidate_dir", + str(tmp_path / "missing-candidate"), + "--output", + str(output), + "--max_abs_diff", + "1", + "--mean_abs_diff", + "1", + "--p99_abs_diff", + "1", + ], + ) + + assert main() == 1 + assert json.loads(output.read_text())["status"] == "invalid_input" diff --git a/tests/test_glm4_moe_lite.py b/tests/test_glm4_moe_lite.py index 852cc01c..6ba0e5d5 100644 --- a/tests/test_glm4_moe_lite.py +++ b/tests/test_glm4_moe_lite.py @@ -1,12 +1,15 @@ from __future__ import annotations +import os from contextlib import ExitStack, contextmanager from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import Mock, patch import pytest import torch +import torch.distributed as dist import torch.nn.functional as F +from torch import nn from transformers import Glm4MoeLiteConfig from transformers.models.glm4_moe_lite.modeling_glm4_moe_lite import ( apply_rotary_pos_emb_interleave, @@ -18,9 +21,11 @@ initialize_sparse_model, ) from sparsevllm.distributed import ParallelContext, ParallelGroup +from sparsevllm.engine.model_runner import ModelRunner from sparsevllm.layers.rotary_embedding import apply_interleaved_rotary_emb from sparsevllm.models.glm4_moe_lite import ( Glm4MoeLiteAttention, + Glm4MoeLiteDecoderLayer, Glm4MoeLiteForCausalLM, Glm4MoeLiteRouter, Glm4MoeLiteSparseMoeBlock, @@ -77,6 +82,16 @@ def _tp_context(tp_rank: int = 0, tp_size: int = 1) -> ParallelContext: ) +def _ep_context(ep_rank: int = 0, ep_size: int = 1) -> ParallelContext: + ranks = tuple(range(ep_size)) + return ParallelContext( + world=ParallelGroup(object(), ranks, ep_rank, ep_size), + tensor=ParallelGroup(None, (ep_rank,), 0, 1), + expert=ParallelGroup(object(), ranks, ep_rank, ep_size), + data=ParallelGroup(None, (ep_rank,), 0, 1), + ) + + def _fake_mla(tp_size: int = 1): spec = MlaAttentionOpSpec( num_q_heads=20, @@ -242,6 +257,29 @@ def test_glm_decode_absorption_and_value_reconstruction_match_linear_algebra() - torch.testing.assert_close(reconstructed, expected_reconstructed) +def test_glm_attention_binds_key_materializer_once_per_manager_context() -> None: + attention = _model().model.layers[0].self_attn + first_manager = SimpleNamespace( + register_attention_key_materializer=Mock() + ) + second_manager = SimpleNamespace( + register_attention_key_materializer=Mock() + ) + + attention._ensure_attention_key_materializer(first_manager, 0) + attention._ensure_attention_key_materializer(first_manager, 0) + attention._ensure_attention_key_materializer(second_manager, 0) + + first_manager.register_attention_key_materializer.assert_called_once_with( + 0, + attention._materialize_attention_keys, + ) + second_manager.register_attention_key_materializer.assert_called_once_with( + 0, + attention._materialize_attention_keys, + ) + + def test_glm_router_uses_bias_only_for_selection_and_scales_weights() -> None: torch.manual_seed(23) router = Glm4MoeLiteRouter(_config()) @@ -308,6 +346,257 @@ def test_tiny_transformers_weights_load_through_strict_glm_mapping() -> None: assert model.model.layers[1].mlp.gate.e_score_correction_bias.dtype == torch.float32 +@pytest.mark.parametrize(("ep_rank", "ep_size"), [(1, 2), (3, 4)]) +def test_glm_ep_loader_keeps_only_local_experts_and_accounts_remote_skips( + ep_rank: int, + ep_size: int, +) -> None: + config = _config() + context = _ep_context(ep_rank, ep_size) + with _construction_context(context): + model = Glm4MoeLiteForCausalLM( + config, + mla_attention=_fake_mla(), + mlp_chunk_size=config.mlp_chunk_size, + decode_cuda_graph=config.decode_cuda_graph, + expect_mtp_weights=False, + ) + initialize_sparse_model(model, config, seed=31) + + experts = model.model.layers[1].mlp.experts + expected_local = 64 // ep_size + assert experts.local_expert_start == ep_rank * expected_local + assert experts.local_expert_end == (ep_rank + 1) * expected_local + assert len(experts._loaded_expert_shards) == expected_local * 3 + assert len(model._intentionally_skipped_expert_weights) == ( + 64 - expected_local + ) * 3 + + +@pytest.mark.parametrize("ep_size", [1, 2, 4]) +def test_glm_sparse_moe_reduces_pure_ep_over_world(ep_size: int) -> None: + context = _ep_context(ep_rank=0, ep_size=ep_size) + block = object.__new__(Glm4MoeLiteSparseMoeBlock) + nn.Module.__init__(block) + block.parallel_context = context + block.mlp_chunk_size = 8 + block.shared_experts = nn.Identity() + block._routed_chunk = lambda hidden_states: hidden_states.clone() + hidden_states = torch.arange(8, dtype=torch.float32).reshape(2, 4) + + with patch.object(dist, "all_reduce") as all_reduce: + output = block(hidden_states) + + torch.testing.assert_close(output, hidden_states * 2) + if ep_size == 1: + all_reduce.assert_not_called() + else: + all_reduce.assert_called_once() + assert all_reduce.call_args.kwargs["group"] is context.world.process_group + + +def test_glm_sparse_moe_reduces_pure_tp_over_world() -> None: + moe_tp_process_group = object() + context = ParallelContext( + world=ParallelGroup(moe_tp_process_group, (0, 1), 0, 2), + tensor=ParallelGroup(moe_tp_process_group, (0, 1), 0, 2), + expert=ParallelGroup(object(), (0,), 0, 1), + data=ParallelGroup(object(), (0,), 0, 1), + moe_tensor=ParallelGroup(moe_tp_process_group, (0, 1), 0, 2), + ) + block = object.__new__(Glm4MoeLiteSparseMoeBlock) + nn.Module.__init__(block) + block.parallel_context = context + block.mlp_chunk_size = 8 + block.shared_experts = nn.Identity() + block._routed_chunk = lambda hidden_states: hidden_states.clone() + hidden_states = torch.arange(8, dtype=torch.float32).reshape(2, 4) + + with patch.object(dist, "all_reduce") as all_reduce: + output = block(hidden_states) + + torch.testing.assert_close(output, hidden_states * 2) + all_reduce.assert_called_once() + assert all_reduce.call_args.kwargs["group"] is context.world.process_group + + +def test_glm_sparse_moe_reduces_hybrid_tp_ep_shards_over_outer_world() -> None: + world_process_group = object() + context = ParallelContext( + world=ParallelGroup(world_process_group, (0, 1, 2, 3), 0, 4), + tensor=ParallelGroup(world_process_group, (0, 1, 2, 3), 0, 4), + expert=ParallelGroup(object(), (0, 2), 0, 2), + data=ParallelGroup(None, (0,), 0, 1), + moe_tensor=ParallelGroup(object(), (0, 1), 0, 2), + ) + block = object.__new__(Glm4MoeLiteSparseMoeBlock) + nn.Module.__init__(block) + block.parallel_context = context + block.mlp_chunk_size = 8 + block.shared_experts = nn.Identity() + block._routed_chunk = lambda hidden_states: hidden_states.clone() + hidden_states = torch.arange(8, dtype=torch.float32).reshape(2, 4) + + with patch.object(dist, "all_reduce") as all_reduce: + output = block(hidden_states) + + torch.testing.assert_close(output, hidden_states * 2) + all_reduce.assert_called_once() + assert all_reduce.call_args.kwargs["group"] is world_process_group + + +def test_glm_moe_debug_contract_populates_model_runner_summaries() -> None: + context = _ep_context() + block = object.__new__(Glm4MoeLiteSparseMoeBlock) + nn.Module.__init__(block) + block.parallel_context = context + block.mlp_chunk_size = 8 + + class _Gate(nn.Module): + def forward(self, hidden_states): + tokens = int(hidden_states.shape[0]) + logits = torch.arange(64, dtype=torch.float32).expand(tokens, -1) + topk_ids = torch.tensor([[1, 3, 5, 7]], dtype=torch.long).expand( + tokens, -1 + ) + topk_weights = torch.full((tokens, 4), 0.25, dtype=torch.float32) + return logits, topk_weights, topk_ids + + class _Experts(nn.Module): + local_expert_start = 0 + local_expert_end = 16 + + def forward(self, hidden_states, topk_ids, topk_weights): + del topk_ids, topk_weights + return hidden_states * 2 + + block.gate = _Gate() + block.experts = _Experts() + block.shared_experts = nn.Identity() + hidden_states = torch.arange(8, dtype=torch.float32).reshape(2, 4) + + with patch.dict(os.environ, {"SPARSEVLLM_DEBUG_MOE": "1"}): + output = block(hidden_states) + + torch.testing.assert_close(output, hidden_states * 3) + torch.testing.assert_close(block.debug_last_local_output, hidden_states * 2) + torch.testing.assert_close(block.debug_last_routed_output, hidden_states * 2) + torch.testing.assert_close(block.debug_last_output, output) + + runner = object.__new__(ModelRunner) + runner.model = SimpleNamespace( + model=SimpleNamespace( + layers=[ + SimpleNamespace(mlp=nn.Identity()), + SimpleNamespace(mlp=block), + ] + ) + ) + runner.parallel_context = context + runner.sparse_controller = SimpleNamespace(debug_state_summary=lambda: {}) + runner.prefix_cache_coordinator = None + runner.debug_last_logits = torch.ones((1, 8), dtype=torch.float32) + runner.rank = 0 + runner.world_size = 1 + runner.device = torch.device("cpu") + + summary = runner.debug_sparse_state_summary() + assert set(summary["moe_synced"]) == {"1"} + assert set(summary["moe_local"]) == {"1"} + assert summary["moe_local"]["1"]["local_expert_start"] == 0 + assert summary["moe_local"]["1"]["local_expert_end"] == 16 + assert summary["moe_local"]["1"]["local_hit_count"] == 8 + assert summary["moe_synced"]["1"]["output"]["shape"] == [2, 4] + assert summary["parallel"]["configured"] == { + "tensor_parallel_size": 1, + "expert_parallel_size": 1, + "data_parallel_size": 1, + "world_size": 1, + } + assert summary["parallel"]["effective"]["expert"] == { + "rank": 0, + "size": 1, + "ranks": [0], + } + assert summary["parallel"]["attention_replicated_for_ep"] is False + + cpu_states = runner.debug_moe_states_cpu() + assert cpu_states is not None + assert set(cpu_states) == {1} + torch.testing.assert_close(cpu_states[1]["output"], output) + consistency = runner.debug_replica_consistency() + assert consistency is not None + assert set(consistency["moe_layers"]) == {"1"} + assert consistency["moe_layers"]["1"]["topk_ids_mismatch"] is False + + +@pytest.mark.parametrize("ep_size", [1, 2, 4]) +def test_glm_decoder_syncs_replicated_attention_before_post_norm( + ep_size: int, +) -> None: + calls: list[str] = [] + context = _ep_context(ep_rank=0, ep_size=ep_size) + layer = object.__new__(Glm4MoeLiteDecoderLayer) + nn.Module.__init__(layer) + layer.parallel_context = context + + class _InputNorm(nn.Module): + def forward(self, hidden_states, residual): + calls.append("input_norm") + return hidden_states + 1, residual + + class _Attention(nn.Module): + def forward(self, positions, hidden_states, rotary_emb): + del positions, rotary_emb + calls.append("attention") + return hidden_states + 2 + + class _PostNorm(nn.Module): + def forward(self, hidden_states, residual): + calls.append("post_norm") + return hidden_states + 3, residual + + class _Mlp(nn.Module): + def forward(self, hidden_states): + calls.append("mlp") + return hidden_states + 4 + + layer.input_layernorm = _InputNorm() + layer.self_attn = _Attention() + layer.post_attention_layernorm = _PostNorm() + layer.mlp = _Mlp() + hidden_states = torch.zeros((1, 4)) + residual = torch.ones((1, 4)) + + def record_broadcast(*args, **kwargs): + calls.append("broadcast") + + with patch.object(dist, "broadcast", side_effect=record_broadcast) as broadcast: + output, actual_residual = layer( + torch.zeros((1,), dtype=torch.long), + hidden_states, + residual, + object(), + ) + + torch.testing.assert_close(output, torch.full_like(output, 10)) + assert actual_residual is residual + if ep_size == 1: + assert calls == ["input_norm", "attention", "post_norm", "mlp"] + broadcast.assert_not_called() + else: + assert calls == [ + "input_norm", + "attention", + "broadcast", + "post_norm", + "mlp", + ] + broadcast.assert_called_once() + assert broadcast.call_args.kwargs["src"] == context.expert.ranks[0] + assert broadcast.call_args.kwargs["group"] is context.expert.process_group + + def test_glm_mtp_skip_set_is_exact() -> None: config = _config() context = _tp_context() diff --git a/tests/test_glm_cuda_graph.py b/tests/test_glm_cuda_graph.py new file mode 100644 index 00000000..8414b23b --- /dev/null +++ b/tests/test_glm_cuda_graph.py @@ -0,0 +1,1363 @@ +from __future__ import annotations + +import json +import os +from collections import deque +from types import SimpleNamespace +from unittest.mock import patch + +import numpy as np +import pytest +import torch +from torch import nn + +from sparsevllm.config import RuntimeLayout +from sparsevllm.configs.model import resolve_attention_qk_head_dim +from sparsevllm.distributed import ParallelContext +from sparsevllm.engine.cache_manager import LayerBatchStates +from sparsevllm.engine.cache_manager.h2o import H2OCacheManager +from sparsevllm.engine.cache_manager.omnikv import OmniKVCacheManager +from sparsevllm.engine.cache_manager.rkv import RKVCacheManager +from sparsevllm.engine.cache_manager.snapkv import SnapKVCacheManager +from sparsevllm.engine.cache_manager.standard import StandardCacheManager +from sparsevllm.engine.cache_manager.storage import MlaLatentStorage +from sparsevllm.engine.cache_manager.streamingllm import ( + StreamingLLMCacheManager, +) +from sparsevllm.engine.decode_cuda_graph import DecodeCudaGraphRunner +from sparsevllm.engine.runtime_state import RuntimeState +from sparsevllm.engine.sequence import Sequence +from sparsevllm.engine.sparse_controller import SparseController +from sparsevllm.layers.mla_attention import MLAAttention +from sparsevllm.layers.rotary_embedding import RotaryEmbedding +from sparsevllm.models.glm4_moe_lite import ( + Glm4MoeLiteAttention, + Glm4MoeLiteForCausalLM, + Glm4MoeLiteSparseMoeBlock, +) +from sparsevllm.operators.mla_attention import MlaAttentionOpSpec +from sparsevllm.utils.context import get_context + +from glm_test_helpers import ( + _glm_hf_config, + _single_rank_parallel_context, + _tensor_sha256, +) + + +def _make_glm_graph_lane( + *, + device: torch.device, + parallel_context: ParallelContext, + attention_state: dict[str, torch.Tensor] | None, + embedding_state: dict[str, torch.Tensor] | None, + head_state: dict[str, torch.Tensor] | None, + initial_latent: torch.Tensor, + initial_rope: torch.Tensor, +): + hf_config = _glm_hf_config( + num_hidden_layers=1, + max_position_embeddings=128, + ) + hf_config.rms_norm_eps = 1e-6 + spec = MlaAttentionOpSpec( + num_q_heads=20, + kv_lora_rank=512, + rope_dim=64, + qk_head_dim=256, + value_head_dim=256, + activation_dtype=torch.bfloat16, + cache_dtype=torch.bfloat16, + tp_size=1, + cuda_graph=True, + ) + mla_attention = MLAAttention.bind( + spec=spec, + device=device, + max_batch_size=1, + prefill_workspace_bytes=1024 * 1024, + hidden_size=64, + projection_chunk_size=8, + ) + previous_dtype = torch.get_default_dtype() + torch.set_default_dtype(torch.bfloat16) + try: + with ( + patch( + "sparsevllm.models.glm4_moe_lite.get_parallel_context", + return_value=parallel_context, + ), + patch( + "sparsevllm.layers.linear.get_parallel_context", + return_value=parallel_context, + ), + torch.device(device), + ): + attention = Glm4MoeLiteAttention( + hf_config, + mla_attention, + projection_chunk_size=8, + ) + embedding = nn.Embedding(128, 64) + lm_head = nn.Linear(64, 128, bias=False) + rotary = RotaryEmbedding( + 64, + 64, + 128, + 1_000_000.0, + backend="torch", + interleaved=True, + ) + finally: + torch.set_default_dtype(previous_dtype) + + if attention_state is None: + generator = torch.Generator(device=device).manual_seed(941) + with torch.no_grad(): + for parameter in attention.parameters(): + parameter.copy_( + torch.randn( + parameter.shape, + dtype=parameter.dtype, + device=device, + generator=generator, + ) + * 0.02 + ) + embedding.weight.copy_( + torch.randn( + embedding.weight.shape, + dtype=embedding.weight.dtype, + device=device, + generator=generator, + ) + * 0.02 + ) + lm_head.weight.copy_( + torch.randn( + lm_head.weight.shape, + dtype=lm_head.weight.dtype, + device=device, + generator=generator, + ) + * 0.02 + ) + else: + assert embedding_state is not None and head_state is not None + attention.load_state_dict(attention_state) + embedding.load_state_dict(embedding_state) + lm_head.load_state_dict(head_state) + + storage = MlaLatentStorage( + kv_lora_rank=512, + rope_dim=64, + dtype=torch.bfloat16, + ) + storage.allocate(num_layers=1, num_slots=16, device=device) + assert storage.latent_cache is not None and storage.rope_cache is not None + storage.latent_cache.zero_() + storage.rope_cache.zero_() + storage.latent_cache[0, :3].copy_(initial_latent) + storage.rope_cache[0, :3].copy_(initial_rope) + + runtime_config = SimpleNamespace( + vllm_sparse_method="", + runtime_layout=RuntimeLayout.dense(1), + hf_config=SimpleNamespace( + num_hidden_layers=1, + num_attention_heads=20, + hidden_size=64, + head_dim=256, + torch_dtype=torch.bfloat16, + ), + obs_layer_ids=[], + full_attn_layers=[], + num_sink_tokens=0, + num_recent_tokens=0, + decode_keep_tokens=0, + sparse_attn_score_dtype="float32", + tensor_parallel_size=1, + decode_cuda_graph=True, + decode_cuda_graph_context_policy="current", + decode_cuda_graph_max_cached_graphs=None, + ) + manager = object.__new__(StandardCacheManager) + manager.config = runtime_config + manager.parallel_context = parallel_context + manager.device = device + manager.runtime_layout = runtime_config.runtime_layout + manager.num_layers = 1 + manager.num_kv_layers = 1 + manager.max_model_len = 128 + manager.attention_cache_storage = storage + manager.kv_cache = None + manager._decode_static_max_context_len = None + manager._attention_key_materializers = {} + manager.free_slots_stack = torch.empty(16, dtype=torch.int32, device=device) + manager.free_slots_stack[:13].copy_( + torch.arange(3, 16, dtype=torch.int32, device=device) + ) + manager._num_free_slots = 13 + manager.buffer_req_to_token_slots = torch.zeros( + (1, 128), + dtype=torch.int32, + device=device, + ) + manager.buffer_req_to_token_slots[0, :3].copy_( + torch.arange(3, dtype=torch.int32, device=device) + ) + sequence = Sequence([5, 7, 11, 13]) + sequence.num_prefilled_tokens = sequence.num_prompt_tokens + sequence.temperature = 0.0 + manager.seq_id_to_row = {sequence.seq_id: 0} + manager.free_rows = deque() + manager.row_seq_lens = np.asarray([3], dtype=np.int32) + manager.layer_batch_state = LayerBatchStates() + manager._decode_static_index_buffers = {} + manager.enable_prefix_caching = False + manager.prefix_cache_block_size = 4 + manager.prefix_cache = None + manager.seq_id_to_prefix_blocks = {} + manager.seq_id_to_cached_ranges = {} + manager._scheduler_capacity_snapshot_depth = 0 + manager._scheduler_freeable_block_ids = None + manager.prefix_offload_controller = None + manager._prefix_offload_step_h2d_operations = [] + manager._prefix_write_through_candidates = {} + manager._init_prefix_cache_runtime() + + sparse_controller = SparseController(runtime_config, manager) + runtime_state = RuntimeState(runtime_config, manager) + + def run_model( + input_ids: torch.Tensor, + positions: torch.Tensor, + is_prefill: bool, + ) -> torch.Tensor: + assert not is_prefill + get_context().now_layer_idx = 0 + hidden_states = embedding(input_ids) + hidden_states = attention(positions, hidden_states, rotary) + return lm_head(hidden_states) + + runner = DecodeCudaGraphRunner( + runtime_state=runtime_state, + cache_manager=manager, + recurrent_state_manager=None, + sparse_controller=sparse_controller, + run_model=run_model, + is_long_text_batch=lambda seqs, is_prefill: False, + method="", + rank=0, + capture_sizes=[1], + context_sizes=[128], + ) + return SimpleNamespace( + attention=attention, + embedding=embedding, + lm_head=lm_head, + manager=manager, + storage=storage, + sequence=sequence, + runner=runner, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@torch.inference_mode() +def test_glm_vanilla_decode_cuda_graph_matches_static_eager(): + device = torch.device("cuda") + parallel_context = _single_rank_parallel_context() + generator = torch.Generator(device=device).manual_seed(937) + initial_latent = torch.randn( + (3, 1, 512), + dtype=torch.bfloat16, + device=device, + generator=generator, + ) + initial_rope = torch.randn( + (3, 1, 64), + dtype=torch.bfloat16, + device=device, + generator=generator, + ) + eager = _make_glm_graph_lane( + device=device, + parallel_context=parallel_context, + attention_state=None, + embedding_state=None, + head_state=None, + initial_latent=initial_latent, + initial_rope=initial_rope, + ) + graph = _make_glm_graph_lane( + device=device, + parallel_context=parallel_context, + attention_state=eager.attention.state_dict(), + embedding_state=eager.embedding.state_dict(), + head_state=eager.lm_head.state_dict(), + initial_latent=initial_latent, + initial_rope=initial_rope, + ) + + step_evidence = [] + captured_graph = None + for step in range(2): + eager_logits = eager.runner.run_eager_static([eager.sequence]) + graph_logits, graph_token_ids = graph.runner.run( + [graph.sequence], + capture_sampling=True, + ) + torch.cuda.synchronize() + assert eager_logits is not None + assert graph_logits is not None + assert graph_token_ids is not None + torch.testing.assert_close(graph_logits, eager_logits, rtol=0, atol=0) + eager_token_ids = eager_logits.argmax(dim=-1) + torch.testing.assert_close(graph_token_ids, eager_token_ids, rtol=0, atol=0) + + graph_states = [ + state + for state in graph.runner._graphs.values() + if state.graph is not None + ] + assert len(graph_states) == 1 + if captured_graph is None: + captured_graph = graph_states[0].graph + else: + assert graph_states[0].graph is captured_graph + + eager_row_len = int(eager.manager.row_seq_lens[0]) + graph_row_len = int(graph.manager.row_seq_lens[0]) + assert eager_row_len == graph_row_len + eager_slots = eager.manager.buffer_req_to_token_slots[0, :eager_row_len].long() + graph_slots = graph.manager.buffer_req_to_token_slots[0, :graph_row_len].long() + eager_latent = eager.storage.latent_cache[0].index_select(0, eager_slots) + graph_latent = graph.storage.latent_cache[0].index_select(0, graph_slots) + eager_rope = eager.storage.rope_cache[0].index_select(0, eager_slots) + graph_rope = graph.storage.rope_cache[0].index_select(0, graph_slots) + torch.testing.assert_close(graph_latent, eager_latent, rtol=0, atol=0) + torch.testing.assert_close(graph_rope, eager_rope, rtol=0, atol=0) + step_evidence.append( + { + "step": step + 1, + "token": int(graph_token_ids[0].item()), + "logits_sha256": _tensor_sha256(graph_logits), + "latent_sha256": _tensor_sha256(graph_latent), + "rope_sha256": _tensor_sha256(graph_rope), + } + ) + if step == 0: + next_token = int(graph_token_ids[0].item()) + eager.sequence.append_token(next_token) + graph.sequence.append_token(next_token) + + evidence = { + "harness_scope": "tiny_random_attention_component", + "real_checkpoint": False, + "graph_active": captured_graph is not None, + "graph_count": sum( + state.graph is not None for state in graph.runner._graphs.values() + ), + "capture_count": graph.runner.capture_count, + "replay_count": graph.runner.replay_count, + "eager_static_count": graph.runner.eager_static_count, + "force_eager": graph.manager.decode_cuda_graph_force_eager(), + "force_eager_count": graph.runner.force_eager_count, + "fallback": graph.runner.force_eager_count > 0, + "steps": step_evidence, + } + assert evidence["capture_count"] == 1 + assert evidence["replay_count"] == 2 + assert evidence["eager_static_count"] == 0 + assert evidence["force_eager_count"] == 0 + assert evidence["fallback"] is False + print("GLM_VANILLA_CUDA_GRAPH_EVIDENCE=" + json.dumps(evidence, sort_keys=True)) + + +def _make_glm_full_graph_lane( + *, + device: torch.device, + parallel_context: ParallelContext, + model_state: dict[str, torch.Tensor] | None, + initial_latent: torch.Tensor, + initial_rope: torch.Tensor, +): + model_config = _glm_hf_config( + num_hidden_layers=2, + mlp_layer_types=["dense", "sparse"], + max_position_embeddings=129, + moe_intermediate_size=16, + hidden_act="silu", + rms_norm_eps=1e-6, + dtype=torch.bfloat16, + tie_word_embeddings=False, + rope_parameters={ + "rope_type": "default", + "rope_theta": 1_000_000.0, + }, + ) + mla_spec = MlaAttentionOpSpec( + num_q_heads=20, + kv_lora_rank=512, + rope_dim=64, + qk_head_dim=256, + value_head_dim=256, + activation_dtype=torch.bfloat16, + cache_dtype=torch.bfloat16, + tp_size=1, + cuda_graph=True, + ) + mla_attention = MLAAttention.bind( + spec=mla_spec, + device=device, + max_batch_size=1, + prefill_workspace_bytes=1024 * 1024, + hidden_size=64, + projection_chunk_size=8, + ) + previous_dtype = torch.get_default_dtype() + torch.set_default_dtype(torch.bfloat16) + try: + with ( + patch( + "sparsevllm.models.glm4_moe_lite.get_parallel_context", + return_value=parallel_context, + ), + patch( + "sparsevllm.layers.linear.get_parallel_context", + return_value=parallel_context, + ), + patch( + "sparsevllm.layers.embed_head.get_parallel_context", + return_value=parallel_context, + ), + torch.device(device), + ): + model = Glm4MoeLiteForCausalLM( + model_config, + mla_attention=mla_attention, + mlp_chunk_size=8, + decode_cuda_graph=True, + expect_mtp_weights=False, + ) + finally: + torch.set_default_dtype(previous_dtype) + + if model_state is None: + generator = torch.Generator(device=device).manual_seed(953) + with torch.no_grad(): + for name, parameter in model.named_parameters(): + if "norm.weight" in name or "layernorm.weight" in name: + parameter.fill_(1.0) + else: + parameter.copy_( + torch.randn( + parameter.shape, + dtype=parameter.dtype, + device=device, + generator=generator, + ) + * 0.02 + ) + else: + model.load_state_dict(model_state) + + num_layers = 2 + storage = MlaLatentStorage( + kv_lora_rank=512, + rope_dim=64, + dtype=torch.bfloat16, + ) + storage.allocate(num_layers=num_layers, num_slots=16, device=device) + assert storage.latent_cache is not None and storage.rope_cache is not None + storage.latent_cache.zero_() + storage.rope_cache.zero_() + storage.latent_cache[:, :3].copy_(initial_latent) + storage.rope_cache[:, :3].copy_(initial_rope) + + runtime_layout = RuntimeLayout.dense(num_layers) + runtime_config = SimpleNamespace( + vllm_sparse_method="", + runtime_layout=runtime_layout, + hf_config=model_config, + obs_layer_ids=[], + full_attn_layers=[], + num_sink_tokens=0, + num_recent_tokens=0, + decode_keep_tokens=0, + sparse_attn_score_dtype="float32", + tensor_parallel_size=1, + decode_cuda_graph=True, + decode_cuda_graph_context_policy="current", + decode_cuda_graph_max_cached_graphs=None, + ) + manager = object.__new__(StandardCacheManager) + manager.config = runtime_config + manager.parallel_context = parallel_context + manager.device = device + manager.runtime_layout = runtime_layout + manager.num_layers = num_layers + manager.num_kv_layers = num_layers + manager.max_model_len = 128 + manager.attention_cache_storage = storage + manager.kv_cache = None + manager._decode_static_max_context_len = None + manager._attention_key_materializers = {} + manager.free_slots_stack = torch.empty(16, dtype=torch.int32, device=device) + manager.free_slots_stack[:13].copy_( + torch.arange(3, 16, dtype=torch.int32, device=device) + ) + manager._num_free_slots = 13 + manager.buffer_req_to_token_slots = torch.zeros( + (1, 128), + dtype=torch.int32, + device=device, + ) + manager.buffer_req_to_token_slots[0, :3].copy_( + torch.arange(3, dtype=torch.int32, device=device) + ) + sequence = Sequence([17, 19, 23, 29]) + sequence.num_prefilled_tokens = sequence.num_prompt_tokens + sequence.temperature = 0.0 + manager.seq_id_to_row = {sequence.seq_id: 0} + manager.free_rows = deque() + manager.row_seq_lens = np.asarray([3], dtype=np.int32) + manager.layer_batch_state = LayerBatchStates() + manager._decode_static_index_buffers = {} + manager.enable_prefix_caching = False + manager.prefix_cache_block_size = 4 + manager.prefix_cache = None + manager.seq_id_to_prefix_blocks = {} + manager.seq_id_to_cached_ranges = {} + manager._scheduler_capacity_snapshot_depth = 0 + manager._scheduler_freeable_block_ids = None + manager.prefix_offload_controller = None + manager._prefix_offload_step_h2d_operations = [] + manager._prefix_write_through_candidates = {} + manager._init_prefix_cache_runtime() + + sparse_controller = SparseController(runtime_config, manager) + model.model.sparse_controller = sparse_controller + runtime_state = RuntimeState(runtime_config, manager) + + def run_model( + input_ids: torch.Tensor, + positions: torch.Tensor, + is_prefill: bool, + ) -> torch.Tensor: + assert not is_prefill + return model.compute_logits(model(input_ids, positions)) + + runner = DecodeCudaGraphRunner( + runtime_state=runtime_state, + cache_manager=manager, + recurrent_state_manager=None, + sparse_controller=sparse_controller, + run_model=run_model, + is_long_text_batch=lambda seqs, is_prefill: False, + method="", + rank=0, + capture_sizes=[1], + context_sizes=[128], + ) + return SimpleNamespace( + model=model, + manager=manager, + storage=storage, + sequence=sequence, + runner=runner, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@torch.inference_mode() +def test_glm_full_decoder_moe_cuda_graph_matches_static_eager(): + device = torch.device("cuda") + parallel_context = _single_rank_parallel_context() + generator = torch.Generator(device=device).manual_seed(947) + initial_latent = torch.randn( + (2, 3, 1, 512), + dtype=torch.bfloat16, + device=device, + generator=generator, + ) + initial_rope = torch.randn( + (2, 3, 1, 64), + dtype=torch.bfloat16, + device=device, + generator=generator, + ) + eager = _make_glm_full_graph_lane( + device=device, + parallel_context=parallel_context, + model_state=None, + initial_latent=initial_latent, + initial_rope=initial_rope, + ) + graph = _make_glm_full_graph_lane( + device=device, + parallel_context=parallel_context, + model_state=eager.model.state_dict(), + initial_latent=initial_latent, + initial_rope=initial_rope, + ) + eager_moe = eager.model.model.layers[1].mlp + graph_moe = graph.model.model.layers[1].mlp + assert isinstance(eager_moe, Glm4MoeLiteSparseMoeBlock) + assert isinstance(graph_moe, Glm4MoeLiteSparseMoeBlock) + assert graph_moe.experts.provider.name == "triton" + + steps = [] + captured_graph = None + with patch.dict(os.environ, {"SPARSEVLLM_DEBUG_MOE": "1"}): + for step in range(2): + eager_logits = eager.runner.run_eager_static([eager.sequence]) + graph_logits, graph_token_ids = graph.runner.run( + [graph.sequence], + capture_sampling=True, + ) + torch.cuda.synchronize() + assert eager_logits is not None + assert graph_logits is not None + assert graph_token_ids is not None + torch.testing.assert_close(graph_logits, eager_logits, rtol=0, atol=0) + eager_token_ids = eager_logits.argmax(dim=-1) + torch.testing.assert_close( + graph_token_ids, + eager_token_ids, + rtol=0, + atol=0, + ) + torch.testing.assert_close( + graph_moe.debug_last_topk_ids, + eager_moe.debug_last_topk_ids, + rtol=0, + atol=0, + ) + torch.testing.assert_close( + graph_moe.debug_last_topk_weights, + eager_moe.debug_last_topk_weights, + rtol=0, + atol=0, + ) + local_hit_count = graph_moe.debug_last_local_hit_count + if isinstance(local_hit_count, torch.Tensor): + local_hit_count = int(local_hit_count.item()) + assert local_hit_count == 4 + + graph_states = [ + state + for state in graph.runner._graphs.values() + if state.graph is not None + ] + assert len(graph_states) == 1 + if captured_graph is None: + captured_graph = graph_states[0].graph + else: + assert graph_states[0].graph is captured_graph + + row_len = int(graph.manager.row_seq_lens[0]) + assert row_len == int(eager.manager.row_seq_lens[0]) + eager_slots = eager.manager.buffer_req_to_token_slots[0, :row_len].long() + graph_slots = graph.manager.buffer_req_to_token_slots[0, :row_len].long() + eager_latent = eager.storage.latent_cache.index_select(1, eager_slots) + graph_latent = graph.storage.latent_cache.index_select(1, graph_slots) + eager_rope = eager.storage.rope_cache.index_select(1, eager_slots) + graph_rope = graph.storage.rope_cache.index_select(1, graph_slots) + torch.testing.assert_close(graph_latent, eager_latent, rtol=0, atol=0) + torch.testing.assert_close(graph_rope, eager_rope, rtol=0, atol=0) + steps.append( + { + "step": step + 1, + "token": int(graph_token_ids[0].item()), + "logits_sha256": _tensor_sha256(graph_logits), + "latent_sha256": _tensor_sha256(graph_latent), + "rope_sha256": _tensor_sha256(graph_rope), + "topk_ids": graph_moe.debug_last_topk_ids.tolist(), + "topk_sha256": _tensor_sha256( + graph_moe.debug_last_topk_ids + ), + "local_expert_hits": local_hit_count, + } + ) + if step == 0: + next_token = int(graph_token_ids[0].item()) + eager.sequence.append_token(next_token) + graph.sequence.append_token(next_token) + + evidence = { + "harness_scope": "tiny_random_full_decoder_moe", + "model_class": "Glm4MoeLiteForCausalLM", + "real_checkpoint": False, + "moe_provider": graph_moe.experts.provider.name, + "graph_active": captured_graph is not None, + "graph_count": sum( + state.graph is not None for state in graph.runner._graphs.values() + ), + "capture_count": graph.runner.capture_count, + "replay_count": graph.runner.replay_count, + "eager_static_count": graph.runner.eager_static_count, + "force_eager_count": graph.runner.force_eager_count, + "fallback": graph.runner.force_eager_count > 0, + "steps": steps, + } + assert evidence["graph_count"] == 1 + assert evidence["capture_count"] == 1 + assert evidence["replay_count"] == 2 + assert evidence["eager_static_count"] == 0 + assert evidence["force_eager_count"] == 0 + assert evidence["fallback"] is False + print("GLM_FULL_CUDA_GRAPH_EVIDENCE=" + json.dumps(evidence, sort_keys=True)) + + +_GLM_GRAPH_METHOD_MANAGERS = { + "streamingllm": StreamingLLMCacheManager, + "snapkv": SnapKVCacheManager, + "h2o": H2OCacheManager, + "omnikv": OmniKVCacheManager, + "rkv": RKVCacheManager, +} + + +def _glm_method_runtime_config(method: str, *, num_layers: int): + hf_config = _glm_hf_config( + num_hidden_layers=num_layers, + max_position_embeddings=128, + ) + hf_config.rms_norm_eps = 1e-6 + return SimpleNamespace( + vllm_sparse_method=method, + runtime_layout=RuntimeLayout.dense(num_layers), + hf_config=hf_config, + obs_layer_ids=[0] if method == "omnikv" else [], + full_attn_layers=[0] if method == "omnikv" else [], + num_sink_tokens=1, + num_recent_tokens=1, + decode_keep_tokens=1, + sparse_attn_score_dtype="float32", + tensor_parallel_size=1, + decode_cuda_graph=True, + decode_cuda_graph_context_policy="current", + decode_cuda_graph_max_cached_graphs=None, + max_model_len=4 if method == "h2o" else 16, + max_num_seqs_in_batch=1, + max_num_seqs_in_gpu=1, + max_num_batched_tokens=16, + chunk_prefill_size=8, + pyramid_layer_ratios=None, + snapkv_num_full_layers=0, + snapkv_window_size=2, + pool_kernel_size=1, + prefill_schedule_policy="chunked", + h2o_decode_budget=3, + h2o_prefill_budget=4, + h2o_recent_ratio=0.5, + h2o_prefill_score_window=2, + rkv_compression_interval=1, + rkv_observation_tokens=1, + rkv_alpha=0.5, + rkv_similarity_threshold=0.0, + rkv_recent_similar_keep=0, + rkv_max_redundancy_tokens=16, + rkv_redundancy_window=0, + enable_prefix_caching=False, + prefix_cache_block_size=4, + ) + + +def _initialize_glm_method_cache_manager( + *, + method: str, + config, + parallel_context: ParallelContext, + device: torch.device, + storage: MlaLatentStorage, + sequence: Sequence, + initial_len: int, +): + manager_type = _GLM_GRAPH_METHOD_MANAGERS[method] + manager = object.__new__(manager_type) + manager.config = config + manager.parallel_context = parallel_context + manager.rank = 0 + manager.world_size = 1 + manager.tp_rank = 0 + manager.tp_size = 1 + manager.ep_rank = 0 + manager.ep_size = 1 + manager.dp_rank = 0 + manager.dp_size = 1 + manager.device = device + manager.hf_config = config.hf_config + manager.head_dim = resolve_attention_qk_head_dim(config.hf_config) + manager.max_model_len = int(config.max_model_len) + manager.max_buffer_rows = 1 + manager.num_layers = int(config.hf_config.num_hidden_layers) + manager.num_kv_layers = manager.num_layers + manager.runtime_layout = config.runtime_layout + manager.attention_cache_storage = storage + manager.kv_cache = None + manager._decode_static_max_context_len = None + manager._attention_key_materializers = {} + num_slots = 16 + free_count = num_slots - int(initial_len) + + if method == "omnikv": + manager.free_slots_stack = torch.empty( + num_slots, + dtype=torch.int32, + device=device, + ) + manager.free_slots_stack[:free_count].copy_( + torch.arange(initial_len, num_slots, dtype=torch.int32, device=device) + ) + manager._num_free_slots = free_count + manager.buffer_req_to_token_slots = torch.zeros( + (1, manager.max_model_len), + dtype=torch.int32, + device=device, + ) + manager.buffer_req_to_token_slots[0, :initial_len].copy_( + torch.arange(initial_len, dtype=torch.int32, device=device) + ) + manager.seq_id_to_row = {sequence.seq_id: 0} + manager.free_rows = deque() + manager.row_seq_lens = np.asarray([initial_len], dtype=np.int32) + manager.layer_batch_state = LayerBatchStates() + manager._decode_static_index_buffers = {} + manager.enable_prefix_caching = False + manager.prefix_cache_block_size = 4 + manager.prefix_cache = None + manager.seq_id_to_prefix_blocks = {} + manager.seq_id_to_cached_ranges = {} + manager._scheduler_capacity_snapshot_depth = 0 + manager._scheduler_freeable_block_ids = None + manager.prefix_offload_controller = None + manager._prefix_offload_step_h2d_operations = [] + manager._prefix_write_through_candidates = {} + manager._init_prefix_cache_runtime() + return manager + + manager.layer_num_slots = [num_slots] * manager.num_layers + manager.free_slots_stack_tensor = torch.empty( + (manager.num_layers, num_slots), + dtype=torch.int32, + device=device, + ) + for layer_idx in range(manager.num_layers): + manager.free_slots_stack_tensor[layer_idx, :free_count].copy_( + torch.arange(initial_len, num_slots, dtype=torch.int32, device=device) + ) + manager.free_slots_stack = [ + manager.free_slots_stack_tensor[layer_idx] + for layer_idx in range(manager.num_layers) + ] + manager._num_free_slots = [free_count] * manager.num_layers + manager.buffer_req_to_token_slots_tensor = torch.zeros( + (manager.num_layers, 1, manager.max_model_len), + dtype=torch.int32, + device=device, + ) + manager.buffer_req_to_token_slots_tensor[:, 0, :initial_len].copy_( + torch.arange(initial_len, dtype=torch.int32, device=device).expand( + manager.num_layers, + -1, + ) + ) + manager.buffer_req_to_token_slots = [ + manager.buffer_req_to_token_slots_tensor[layer_idx] + for layer_idx in range(manager.num_layers) + ] + manager.seq_id_to_row = [ + {sequence.seq_id: 0} for _ in range(manager.num_layers) + ] + manager.free_rows = [deque() for _ in range(manager.num_layers)] + manager.row_seq_lens = [ + np.asarray([initial_len], dtype=np.int32) + for _ in range(manager.num_layers) + ] + manager.layer_batch_states = [ + LayerBatchStates() for _ in range(manager.num_layers) + ] + manager._decode_static_buffers = {} + manager._decode_static_index_buffers = {} + manager._decode_static_state_binding_key = None + manager._prefill_attn_score_accumulators = {} + manager._prefill_score_bounds = None + manager._uniform_decode_metadata = method == "streamingllm" + manager.pyramidkv_prefill_staging_num_slots = 0 + manager.pyramidkv_prefill_staging_kv_cache = None + manager._pyramidkv_prefill_staging_active = False + manager._pyramidkv_prefill_staging_was_active = False + manager._pyramidkv_prefill_staging_slot_mapping = None + manager._pyramidkv_prefill_staging_active_slots = None + manager._pyramidkv_prefill_staging_req_indices = None + manager._pyramidkv_prefill_staging_context_lens = None + manager._pyramidkv_prefill_staging_seq_offsets = {} + manager._pyramidkv_prefill_staging_materialized_layers = set() + manager._pyramidkv_long_prefill_offload_step_active = False + manager._pyramidkv_long_prefill_offload_seq_id = None + manager._pyramidkv_long_prefill_offload_start = 0 + manager._pyramidkv_long_prefill_offload_end = 0 + manager._pyramidkv_long_prefill_offload_total_len = 0 + manager._pyramidkv_long_prefill_offload_is_last_chunk = False + manager._pyramidkv_long_prefill_offload_prefetch_stream = None + manager._pyramidkv_long_prefill_offload_prefetch_states = {} + manager.raw_kv_offload_buffer = SimpleNamespace( + release_layer=lambda **_kwargs: None, + ) + + if method == "h2o": + manager._h2o_scores = { + (layer_idx, sequence.seq_id): torch.zeros( + initial_len, + dtype=torch.float32, + device=device, + ) + for layer_idx in range(manager.num_layers) + } + manager._h2o_recent_cursors = {} + manager._h2o_counters = { + "intermediate_prefill_evictions": 0, + "final_prefill_evictions": 0, + "decode_evictions": 0, + "dropped_tokens": 0, + } + manager._h2o_ring_counters = {"fast_rows": 0, "fallback_rows": 0} + manager._h2o_final_prefill_workspace = None + manager._h2o_decode_static_rows = None + manager._h2o_decode_static_topology = None + + if method == "rkv": + obs = int(config.rkv_observation_tokens) + heads = int(config.hf_config.num_attention_heads) + manager._rkv_query_cache_enabled = True + manager._rkv_observation_tokens = obs + manager._rkv_vectorized_prefill_query_cache = True + manager._rkv_batch_clear_query_cache_rows = True + manager._rkv_query_score_static_buffers = {} + manager._rkv_query_cache = [ + torch.zeros( + (1, obs, heads, 256), + dtype=torch.bfloat16, + device=device, + ) + for _ in range(manager.num_layers) + ] + manager._rkv_query_positions = [ + torch.full( + (1, obs), + -1, + dtype=torch.int32, + device=device, + ) + for _ in range(manager.num_layers) + ] + return manager + + +def _make_glm_method_graph_lane( + *, + method: str, + device: torch.device, + parallel_context: ParallelContext, + attention_states: list[dict[str, torch.Tensor]] | None, + embedding_state: dict[str, torch.Tensor] | None, + head_state: dict[str, torch.Tensor] | None, + initial_latent: torch.Tensor, + initial_rope: torch.Tensor, +): + num_layers = 2 + config = _glm_method_runtime_config(method, num_layers=num_layers) + spec = MlaAttentionOpSpec( + num_q_heads=20, + kv_lora_rank=512, + rope_dim=64, + qk_head_dim=256, + value_head_dim=256, + activation_dtype=torch.bfloat16, + cache_dtype=torch.bfloat16, + tp_size=1, + cuda_graph=True, + ) + mla_attention = MLAAttention.bind( + spec=spec, + device=device, + max_batch_size=1, + prefill_workspace_bytes=1024 * 1024, + hidden_size=64, + projection_chunk_size=8, + ) + previous_dtype = torch.get_default_dtype() + torch.set_default_dtype(torch.bfloat16) + try: + with ( + patch( + "sparsevllm.models.glm4_moe_lite.get_parallel_context", + return_value=parallel_context, + ), + patch( + "sparsevllm.layers.linear.get_parallel_context", + return_value=parallel_context, + ), + torch.device(device), + ): + attentions = nn.ModuleList( + [ + Glm4MoeLiteAttention( + config.hf_config, + mla_attention, + projection_chunk_size=8, + ) + for _ in range(num_layers) + ] + ) + embedding = nn.Embedding(128, 64) + lm_head = nn.Linear(64, 128, bias=False) + rotary = RotaryEmbedding( + 64, + 64, + 128, + 1_000_000.0, + backend="torch", + interleaved=True, + ) + finally: + torch.set_default_dtype(previous_dtype) + + if attention_states is None: + generator = torch.Generator(device=device).manual_seed(967) + with torch.no_grad(): + for parameter in attentions.parameters(): + parameter.copy_( + torch.randn( + parameter.shape, + dtype=parameter.dtype, + device=device, + generator=generator, + ) + * 0.02 + ) + embedding.weight.copy_( + torch.randn( + embedding.weight.shape, + dtype=embedding.weight.dtype, + device=device, + generator=generator, + ) + * 0.02 + ) + lm_head.weight.copy_( + torch.randn( + lm_head.weight.shape, + dtype=lm_head.weight.dtype, + device=device, + generator=generator, + ) + * 0.02 + ) + else: + assert embedding_state is not None and head_state is not None + assert len(attention_states) == num_layers + for attention, state in zip(attentions, attention_states): + attention.load_state_dict(state) + embedding.load_state_dict(embedding_state) + lm_head.load_state_dict(head_state) + + storage = MlaLatentStorage( + kv_lora_rank=512, + rope_dim=64, + dtype=torch.bfloat16, + ) + storage.allocate(num_layers=num_layers, num_slots=16, device=device) + assert storage.latent_cache is not None and storage.rope_cache is not None + storage.latent_cache.zero_() + storage.rope_cache.zero_() + initial_len = int(initial_latent.shape[1]) + storage.latent_cache[:, :initial_len].copy_(initial_latent) + storage.rope_cache[:, :initial_len].copy_(initial_rope) + + sequence = Sequence([5, 7, 11, 13]) + sequence.num_prefilled_tokens = sequence.num_prompt_tokens + sequence.temperature = 0.0 + sequence.max_tokens = 8 + manager = _initialize_glm_method_cache_manager( + method=method, + config=config, + parallel_context=parallel_context, + device=device, + storage=storage, + sequence=sequence, + initial_len=initial_len, + ) + controller = SparseController(config, manager) + runtime_state = RuntimeState(config, manager) + + def run_model( + input_ids: torch.Tensor, + positions: torch.Tensor, + is_prefill: bool, + ) -> torch.Tensor: + assert not is_prefill + context = get_context() + hidden_states = embedding(input_ids) + for layer_idx, attention in enumerate(attentions): + context.now_layer_idx = layer_idx + hidden_states = attention(positions, hidden_states, rotary) + controller.on_layer_end(layer_idx, context) + return lm_head(hidden_states) + + runner = DecodeCudaGraphRunner( + runtime_state=runtime_state, + cache_manager=manager, + recurrent_state_manager=None, + sparse_controller=controller, + run_model=run_model, + is_long_text_batch=lambda seqs, is_prefill: True, + method=method, + rank=0, + capture_sizes=[1], + context_sizes=[16], + ) + return SimpleNamespace( + attentions=attentions, + embedding=embedding, + lm_head=lm_head, + manager=manager, + storage=storage, + sequence=sequence, + controller=controller, + runtime_state=runtime_state, + runner=runner, + ) + + +def _glm_method_row_lens(lane, method: str) -> list[int]: + if method == "omnikv": + return [int(lane.manager.row_seq_lens[0])] * len(lane.attentions) + return [ + int(lane.manager.row_seq_lens[layer_idx][0]) + for layer_idx in range(len(lane.attentions)) + ] + + +def _glm_method_slot_rows(lane, method: str) -> list[list[int]]: + row_lens = _glm_method_row_lens(lane, method) + if method == "omnikv": + row = lane.manager.buffer_req_to_token_slots[0, : row_lens[0]] + return [row.tolist() for _ in lane.attentions] + return [ + lane.manager.buffer_req_to_token_slots[layer_idx][ + 0, : row_lens[layer_idx] + ].tolist() + for layer_idx in range(len(lane.attentions)) + ] + + +def _glm_method_trigger_state(lane, method: str) -> dict[str, object]: + state: dict[str, object] = { + "row_lens": _glm_method_row_lens(lane, method), + } + if method in {"snapkv", "h2o"}: + scores = lane.controller.layer_batch_sparse_states[0].attn_score + assert scores is not None + state.update( + score_ptr=int(scores.data_ptr()), + score_sha256=_tensor_sha256(scores), + written_score_count=int((scores > -1e19).sum().item()), + ) + elif method == "omnikv": + target = lane.controller.layer_batch_sparse_states[1] + assert target.active_indices is not None + assert target.active_slots is not None + assert target.context_lens is not None + state.update( + selection_ptr=int(target.active_slots.data_ptr()), + active_indices=target.active_indices.tolist(), + active_slots=target.active_slots.tolist(), + active_context_lens=target.context_lens.tolist(), + ) + elif method == "rkv": + positions = lane.manager._rkv_query_positions[0] + state.update( + query_cache_ptr=int(lane.manager._rkv_query_cache[0].data_ptr()), + query_positions=positions.tolist(), + materializer_bound=lane.manager.has_attention_key_materializer(0), + ) + else: + mapping = lane.manager.layer_batch_states[0].slot_mapping + assert mapping is not None + state["mapping_ptr"] = int(mapping.data_ptr()) + return state + + +@pytest.mark.parametrize( + "method", + ["streamingllm", "snapkv", "h2o", "omnikv", "rkv"], +) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@torch.inference_mode() +def test_glm_sparse_method_decode_cuda_graph_triggers_runtime_path(method: str): + device = torch.device("cuda", torch.cuda.current_device()) + parallel_context = _single_rank_parallel_context() + num_layers = 2 + generator = torch.Generator(device=device).manual_seed(971) + initial_latent = torch.randn( + (num_layers, 3, 1, 512), + dtype=torch.bfloat16, + device=device, + generator=generator, + ) + initial_rope = torch.randn( + (num_layers, 3, 1, 64), + dtype=torch.bfloat16, + device=device, + generator=generator, + ) + eager = _make_glm_method_graph_lane( + method=method, + device=device, + parallel_context=parallel_context, + attention_states=None, + embedding_state=None, + head_state=None, + initial_latent=initial_latent, + initial_rope=initial_rope, + ) + graph = _make_glm_method_graph_lane( + method=method, + device=device, + parallel_context=parallel_context, + attention_states=[ + attention.state_dict() for attention in eager.attentions + ], + embedding_state=eager.embedding.state_dict(), + head_state=eager.lm_head.state_dict(), + initial_latent=initial_latent, + initial_rope=initial_rope, + ) + + steps = [] + stable_ptr = None + trigger_count = 0 + for step in range(2): + eager_logits = eager.runner.run_eager_static([eager.sequence]) + eager_before = _glm_method_trigger_state(eager, method) + eager.controller.post_forward([eager.sequence], is_prefill=False) + eager.runtime_state.on_forward_end([eager.sequence], is_prefill=False) + eager_after_lens = _glm_method_row_lens(eager, method) + + graph_logits, graph_token_ids = graph.runner.run( + [graph.sequence], + capture_sampling=True, + ) + torch.cuda.synchronize() + graph_before = _glm_method_trigger_state(graph, method) + graph.controller.post_forward([graph.sequence], is_prefill=False) + graph.runtime_state.on_forward_end([graph.sequence], is_prefill=False) + graph_after_lens = _glm_method_row_lens(graph, method) + + assert eager_logits is not None + assert graph_logits is not None + assert graph_token_ids is not None + torch.testing.assert_close(graph_logits, eager_logits, rtol=0, atol=0) + eager_token_ids = eager_logits.argmax(dim=-1) + torch.testing.assert_close(graph_token_ids, eager_token_ids, rtol=0, atol=0) + assert graph_before["row_lens"] == eager_before["row_lens"] + assert graph_after_lens == eager_after_lens + assert _glm_method_slot_rows(graph, method) == _glm_method_slot_rows( + eager, + method, + ) + torch.testing.assert_close( + graph.storage.latent_cache, + eager.storage.latent_cache, + rtol=0, + atol=0, + ) + torch.testing.assert_close( + graph.storage.rope_cache, + eager.storage.rope_cache, + rtol=0, + atol=0, + ) + + if method in {"snapkv", "h2o"}: + assert graph_before["score_sha256"] == eager_before["score_sha256"] + pointer = int(graph_before["score_ptr"]) + assert int(graph_before["written_score_count"]) > 0 + elif method == "omnikv": + assert graph_before["active_indices"] == eager_before["active_indices"] + assert graph_before["active_slots"] == eager_before["active_slots"] + assert graph_before["active_context_lens"] == [3] + pointer = int(graph_before["selection_ptr"]) + elif method == "rkv": + assert graph_before["query_positions"] == eager_before["query_positions"] + assert graph_before["materializer_bound"] is True + assert graph_before["query_positions"] == [[3]] + pointer = int(graph_before["query_cache_ptr"]) + else: + pointer = int(graph_before["mapping_ptr"]) + if stable_ptr is None: + stable_ptr = pointer + else: + assert pointer == stable_ptr + + before_len = int(graph_before["row_lens"][0]) + after_len = int(graph_after_lens[0]) + triggered = ( + after_len < before_len + if method != "omnikv" + else int(graph_before["active_context_lens"][0]) < before_len + ) + trigger_count += int(triggered) + if method == "rkv": + assert graph.manager._rkv_query_positions[0].tolist() == [[-1]] + + steps.append( + { + "step": step + 1, + "token": int(graph_token_ids.item()), + "logits_sha256": _tensor_sha256(graph_logits), + "latent_sha256": _tensor_sha256(graph.storage.latent_cache), + "rope_sha256": _tensor_sha256(graph.storage.rope_cache), + "before_lens": graph_before["row_lens"], + "after_lens": graph_after_lens, + "triggered": triggered, + "runtime_state": graph_before, + } + ) + eager.sequence.append_token(int(eager_token_ids.item())) + graph.sequence.append_token(int(graph_token_ids.item())) + + assert trigger_count > 0 + if method == "h2o": + assert graph.manager._h2o_counters["decode_evictions"] == 4 + evidence = { + "method": method, + "harness_scope": "tiny_random_sparse_method_component", + "real_checkpoint": False, + "graph_active": any( + state.graph is not None for state in graph.runner._graphs.values() + ), + "graph_count": sum( + state.graph is not None for state in graph.runner._graphs.values() + ), + "capture_count": graph.runner.capture_count, + "replay_count": graph.runner.replay_count, + "eager_static_count": graph.runner.eager_static_count, + "force_eager_count": graph.runner.force_eager_count, + "fallback": graph.runner.force_eager_count > 0, + "trigger_count": trigger_count, + "stable_runtime_ptr": stable_ptr, + "steps": steps, + } + assert evidence["graph_active"] is True + assert evidence["graph_count"] == 1 + assert evidence["capture_count"] == 1 + assert evidence["replay_count"] == 2 + assert evidence["eager_static_count"] == 0 + assert evidence["force_eager_count"] == 0 + assert evidence["fallback"] is False + print("GLM_METHOD_CUDA_GRAPH_EVIDENCE=" + json.dumps(evidence, sort_keys=True)) diff --git a/tests/test_glm_mla_prefix_cache.py b/tests/test_glm_mla_prefix_cache.py new file mode 100644 index 00000000..f55162cb --- /dev/null +++ b/tests/test_glm_mla_prefix_cache.py @@ -0,0 +1,517 @@ +from __future__ import annotations + +from collections import deque +from types import SimpleNamespace + +import numpy as np +import torch + +from sparsevllm.config import RuntimeLayout +from sparsevllm.engine.cache_manager import ( + AttentionViewMeta, + LayerBatchStates, + PrefillComputeView, +) +from sparsevllm.engine.cache_manager.h2o import H2OCacheManager +from sparsevllm.engine.cache_manager.rkv import RKVCacheManager +from sparsevllm.engine.cache_manager.snapkv import SnapKVCacheManager +from sparsevllm.engine.cache_manager.storage import MlaLatentStorage +from sparsevllm.engine.chain_cache import ChainCacheCoordinator +from sparsevllm.engine.sequence import Sequence +from sparsevllm.engine.sparse_controller import SparseController + + +def _latent_chain_manager(manager_type, method: str): + capacity = 32 + config = SimpleNamespace( + vllm_sparse_method=method, + model="/models/glm-chain-test", + hf_config=SimpleNamespace( + model_type="glm4_moe_lite", + torch_dtype=torch.bfloat16, + num_attention_heads=2, + num_key_value_heads=1, + ), + tensor_parallel_size=1, + max_model_len=capacity, + max_num_seqs_in_gpu=1, + prefix_cache_salt="", + chain_cache_max_tombstones=8, + full_attn_layers=[0], + num_sink_tokens=1, + num_recent_tokens=1, + decode_keep_tokens=2, + snapkv_window_size=2, + snapkv_num_full_layers=0, + sparse_attn_score_dtype="float32", + pool_kernel_size=1, + pyramid_layer_ratios=None, + prefill_schedule_policy="chunked", + chunk_prefill_size=8, + h2o_decode_budget=4, + h2o_prefill_budget=8, + h2o_recent_ratio=0.5, + h2o_prefill_score_window=2, + rkv_compression_interval=2, + rkv_observation_tokens=2, + rkv_alpha=0.5, + rkv_similarity_threshold=0.0, + rkv_recent_similar_keep=0, + rkv_max_redundancy_tokens=16, + rkv_redundancy_window=0, + ) + storage = MlaLatentStorage( + kv_lora_rank=512, + rope_dim=64, + dtype=torch.bfloat16, + ) + storage.allocate(num_layers=1, num_slots=capacity, device=torch.device("cpu")) + + manager = object.__new__(manager_type) + manager.config = config + manager.hf_config = config.hf_config + manager.device = torch.device("cpu") + manager.tp_size = 1 + manager.head_dim = 256 + manager.max_model_len = capacity + manager.max_buffer_rows = 1 + manager.num_layers = 1 + manager.num_kv_layers = 1 + manager.runtime_layout = RuntimeLayout.dense(1) + manager.attention_cache_storage = storage + manager.kv_cache = None + manager._attention_key_materializers = {} + manager.layer_num_slots = [capacity] + manager.free_slots_stack_tensor = torch.arange( + capacity, dtype=torch.int32 + ).view(1, capacity) + manager.free_slots_stack = [manager.free_slots_stack_tensor[0]] + manager._num_free_slots = [capacity] + manager.buffer_req_to_token_slots_tensor = torch.zeros( + (1, 1, capacity), dtype=torch.int32 + ) + manager.buffer_req_to_token_slots = [ + manager.buffer_req_to_token_slots_tensor[0] + ] + manager.seq_id_to_row = [{}] + manager.free_rows = [deque([0])] + manager.row_seq_lens = [np.zeros((1,), dtype=np.int32)] + manager.layer_batch_states = [LayerBatchStates()] + manager._decode_static_state_binding_key = None + manager._decode_static_buffers = {} + manager._decode_static_index_buffers = {} + manager._prefill_attn_score_accumulators = {} + manager._prefill_score_bounds = None + manager._uniform_decode_metadata = False + manager._h2o_scores = {} + manager._h2o_recent_cursors = {} + manager._h2o_counters = { + "intermediate_prefill_evictions": 0, + "final_prefill_evictions": 0, + "decode_evictions": 0, + "dropped_tokens": 0, + } + manager._h2o_ring_counters = { + "fast_rows": 0, + "fallback_rows": 0, + } + manager._h2o_final_prefill_workspace = None + manager._rkv_query_cache_enabled = True + manager._rkv_observation_tokens = 2 + manager._rkv_vectorized_prefill_query_cache = True + manager._rkv_batch_clear_query_cache_rows = True + manager._rkv_query_score_static_buffers = {} + manager._rkv_query_cache = [ + torch.zeros((1, 2, 2, 256), dtype=torch.bfloat16) + ] + manager._rkv_query_positions = [ + torch.full((1, 2), -1, dtype=torch.int32) + ] + manager._pyramidkv_prefill_staging_active = False + manager._pyramidkv_prefill_staging_was_active = False + manager.pyramidkv_prefill_staging_kv_cache = None + manager.pyramidkv_prefill_staging_num_slots = 0 + manager._pyramidkv_long_prefill_offload_prefetch_states = {} + manager.raw_kv_offload_buffer = SimpleNamespace( + release_layer=lambda **_kwargs: None, + ) + return manager, storage, config + + +def _fill_latent_slots(storage, slots: torch.Tensor, values: list[int]) -> None: + payload = storage.layer_payload(0) + values_tensor = torch.tensor(values, dtype=torch.bfloat16) + payload.latent_cache[slots] = values_tensor.view(-1, 1, 1).expand( + -1, 1, 512 + ) + payload.rope_cache[slots] = (values_tensor + 100).view( + -1, 1, 1 + ).expand(-1, 1, 64) + + +def test_snapkv_chain_resume_preserves_latent_payload_and_resets_prefill_scores(): + manager, storage, config = _latent_chain_manager( + SnapKVCacheManager, + "snapkv", + ) + coordinator = ChainCacheCoordinator(config, manager) + owner_tokens = list(range(8)) + owner = Sequence(owner_tokens) + owner.seq_id = 0 + owner.chain_id = "chain-snap-latent" + owner.chain_status = "created" + owner.current_chunk_size = len(owner_tokens) + created = coordinator.plan_admission( + chain_id=owner.chain_id, + seq_id=owner.seq_id, + token_ids=owner_tokens, + ) + assert created.status == "created" + coordinator.apply_admission(created) + + manager._prepare_prefill([owner]) + owner_slots = manager.layer_batch_states[0].slot_mapping.clone().long() + _fill_latent_slots(storage, owner_slots, owner_tokens) + manager._prefill_attn_score_accumulators[(0, owner.seq_id)] = torch.tensor( + [0.0, 1.0, 2.0, 9.0, 3.0, 8.0, 4.0, 0.0] + ) + controller = object.__new__(SparseController) + controller.cache_manager = manager + controller.device = torch.device("cpu") + controller.sparse_method = "snapkv" + controller.num_layers = 1 + controller.num_sink = 1 + controller.num_recent = 1 + controller.decode_keep_tokens = 2 + controller.config = config + controller._is_kv_layer = lambda layer_idx: int(layer_idx) == 0 + controller._kv_layer_index = lambda layer_idx: int(layer_idx) + controller._snapkv_prefill_eviction([owner]) + assert manager._prefill_attn_score_accumulators == {} + resident_slots = manager.buffer_req_to_token_slots[0][0, :4].clone().long() + assert resident_slots.tolist() == owner_slots[[0, 3, 5, 7]].tolist() + payload = storage.layer_payload(0) + torch.testing.assert_close( + payload.latent_cache[resident_slots, 0, 0], + torch.tensor([0, 3, 5, 7], dtype=torch.bfloat16), + ) + torch.testing.assert_close( + payload.rope_cache[resident_slots, 0, 0], + torch.tensor([100, 103, 105, 107], dtype=torch.bfloat16), + ) + coordinator.index.finish( + owner.chain_id, + token_ids=owner_tokens, + processed_token_count=len(owner_tokens), + physical_slots_by_layer=manager.chain_physical_residency(owner.seq_id), + ) + + resumed_tokens = owner_tokens + [8, 9] + resumed_plan = coordinator.plan_admission( + chain_id=owner.chain_id, + seq_id=owner.seq_id, + token_ids=resumed_tokens, + ) + assert resumed_plan.status == "resumed" + assert resumed_plan.reused_tokens == len(owner_tokens) + coordinator.apply_admission(resumed_plan) + resumed = Sequence(resumed_tokens) + resumed.seq_id = owner.seq_id + resumed.chain_id = owner.chain_id + resumed.chain_status = "resumed" + resumed.chain_reused_tokens = len(owner_tokens) + resumed.num_prefilled_tokens = len(owner_tokens) + resumed.current_chunk_size = 2 + manager._prefill_attn_score_accumulators[(0, resumed.seq_id)] = torch.full( + (4,), 999.0 + ) + + input_ids, positions, _ = manager._prepare_prefill([resumed]) + assert input_ids.tolist() == [8, 9] + assert positions.tolist() == [8, 9] + assert manager._prefill_attn_score_accumulators == {} + resumed_row = manager.seq_id_to_row[0][resumed.seq_id] + resumed_slots = manager.buffer_req_to_token_slots[0][ + resumed_row, :6 + ].clone().long() + assert resumed_slots[:4].tolist() == resident_slots.tolist() + _fill_latent_slots(storage, resumed_slots[4:], [8, 9]) + torch.testing.assert_close( + payload.latent_cache[resumed_slots, 0, 0], + torch.tensor([0, 3, 5, 7, 8, 9], dtype=torch.bfloat16), + ) + torch.testing.assert_close( + payload.rope_cache[resumed_slots, 0, 0], + torch.tensor([100, 103, 105, 107, 108, 109], dtype=torch.bfloat16), + ) + + coordinator.index.finish( + owner.chain_id, + token_ids=resumed_tokens, + processed_token_count=len(resumed_tokens), + physical_slots_by_layer=manager.chain_physical_residency(resumed.seq_id), + ) + coordinator.invalidate(owner.chain_id) + manager.free_seq(resumed.seq_id) + assert manager._prefill_attn_score_accumulators == {} + assert manager.seq_id_to_row == [{}] + assert manager._num_free_slots == [32] + + +def test_h2o_chain_resume_preserves_aligned_scores_and_cleans_side_state(): + manager, storage, config = _latent_chain_manager(H2OCacheManager, "h2o") + coordinator = ChainCacheCoordinator(config, manager) + owner_tokens = list(range(6)) + owner = Sequence(owner_tokens) + owner.seq_id = 0 + owner.chain_id = "chain-h2o-latent" + owner.chain_status = "created" + owner.current_chunk_size = len(owner_tokens) + created = coordinator.plan_admission( + chain_id=owner.chain_id, + seq_id=owner.seq_id, + token_ids=owner_tokens, + ) + assert created.status == "created" + coordinator.apply_admission(created) + + manager._prepare_prefill([owner]) + owner_slots = manager.layer_batch_states[0].slot_mapping.clone().long() + _fill_latent_slots(storage, owner_slots, owner_tokens) + manager._h2o_scores[(0, owner.seq_id)] = torch.tensor( + [1.0, 9.0, 2.0, 8.0, 0.0, 0.0] + ) + manager.evict_after_prefill([owner]) + assert manager.row_seq_lens[0].tolist() == [4] + resident_slots = manager.buffer_req_to_token_slots[0][0, :4].clone().long() + payload = storage.layer_payload(0) + torch.testing.assert_close( + payload.latent_cache[resident_slots, 0, 0], + torch.tensor([1, 3, 4, 5], dtype=torch.bfloat16), + ) + torch.testing.assert_close( + payload.rope_cache[resident_slots, 0, 0], + torch.tensor([101, 103, 104, 105], dtype=torch.bfloat16), + ) + assert manager._h2o_scores[(0, owner.seq_id)].tolist() == [9.0, 8.0, 0.0, 0.0] + assert manager._h2o_recent_cursors[(0, owner.seq_id)] == 2 + assert manager._h2o_final_prefill_workspace is None + coordinator.index.finish( + owner.chain_id, + token_ids=owner_tokens, + processed_token_count=len(owner_tokens), + physical_slots_by_layer=manager.chain_physical_residency(owner.seq_id), + ) + + resumed_tokens = owner_tokens + [6, 7] + resumed_plan = coordinator.plan_admission( + chain_id=owner.chain_id, + seq_id=owner.seq_id, + token_ids=resumed_tokens, + ) + assert resumed_plan.status == "resumed" + assert resumed_plan.reused_tokens == len(owner_tokens) + coordinator.apply_admission(resumed_plan) + resumed = Sequence(resumed_tokens) + resumed.seq_id = owner.seq_id + resumed.chain_id = owner.chain_id + resumed.chain_status = "resumed" + resumed.chain_reused_tokens = len(owner_tokens) + resumed.num_prefilled_tokens = len(owner_tokens) + resumed.current_chunk_size = 2 + + input_ids, positions, _ = manager._prepare_prefill([resumed]) + assert input_ids.tolist() == [6, 7] + assert positions.tolist() == [6, 7] + assert manager._h2o_scores[(0, resumed.seq_id)].tolist() == [9.0, 8.0, 0.0, 0.0] + assert manager._h2o_recent_cursors[(0, resumed.seq_id)] == 2 + resumed_row = manager.seq_id_to_row[0][resumed.seq_id] + resumed_slots = manager.buffer_req_to_token_slots[0][ + resumed_row, :6 + ].clone().long() + assert resumed_slots[:4].tolist() == resident_slots.tolist() + _fill_latent_slots(storage, resumed_slots[4:], [6, 7]) + manager._h2o_scores[(0, resumed.seq_id)] = manager._accumulate_score( + manager._h2o_scores[(0, resumed.seq_id)], + torch.tensor([0.0, 0.0, 0.0, 0.0, 7.0, 6.0]), + new_len=6, + weight=1.0, + ) + manager.evict_after_prefill([resumed]) + final_slots = manager.buffer_req_to_token_slots[0][ + resumed_row, :4 + ].clone().long() + torch.testing.assert_close( + payload.latent_cache[final_slots, 0, 0], + torch.tensor([1, 3, 6, 7], dtype=torch.bfloat16), + ) + torch.testing.assert_close( + payload.rope_cache[final_slots, 0, 0], + torch.tensor([101, 103, 106, 107], dtype=torch.bfloat16), + ) + assert manager._h2o_scores[(0, resumed.seq_id)].tolist() == [9.0, 8.0, 7.0, 6.0] + assert manager._h2o_recent_cursors[(0, resumed.seq_id)] == 2 + assert manager._h2o_counters["final_prefill_evictions"] == 2 + assert manager._h2o_final_prefill_workspace is None + + coordinator.index.finish( + owner.chain_id, + token_ids=resumed_tokens, + processed_token_count=len(resumed_tokens), + physical_slots_by_layer=manager.chain_physical_residency(resumed.seq_id), + ) + coordinator.invalidate(owner.chain_id) + manager.free_seq(resumed.seq_id) + assert manager._h2o_scores == {} + assert manager._h2o_recent_cursors == {} + assert manager._prefill_attn_score_accumulators == {} + assert manager.seq_id_to_row == [{}] + assert manager._num_free_slots == [32] + + +def _rkv_prefill_view(manager, storage) -> PrefillComputeView: + state = manager.layer_batch_states[0] + return PrefillComputeView( + meta=AttentionViewMeta( + active_slots=state.slot_mapping, + req_indices=state.req_indices, + context_lens=state.context_lens, + max_context_len=state.max_context_len, + ), + payload=storage.layer_payload(0), + ) + + +def test_rkv_chain_resume_rebuilds_query_observations_without_cross_turn_leak(): + manager, storage, config = _latent_chain_manager(RKVCacheManager, "rkv") + coordinator = ChainCacheCoordinator(config, manager) + owner_tokens = list(range(6)) + owner = Sequence(owner_tokens) + owner.seq_id = 0 + owner.chain_id = "chain-rkv-latent" + owner.chain_status = "created" + owner.current_chunk_size = len(owner_tokens) + created = coordinator.plan_admission( + chain_id=owner.chain_id, + seq_id=owner.seq_id, + token_ids=owner_tokens, + ) + assert created.status == "created" + coordinator.apply_admission(created) + + manager._prepare_prefill([owner]) + owner_slots = manager.layer_batch_states[0].slot_mapping.clone().long() + _fill_latent_slots(storage, owner_slots, owner_tokens) + owner_q = torch.arange(6, dtype=torch.bfloat16).view(6, 1, 1).expand( + 6, 2, 256 + ) + manager.record_prefill_query( + 0, + owner_q, + _rkv_prefill_view(manager, storage), + b_start_loc=torch.tensor([0], dtype=torch.int32), + chunk_lens=torch.tensor([6], dtype=torch.int32), + ) + assert manager._rkv_query_positions[0][0].tolist() == [4, 5] + torch.testing.assert_close( + manager._rkv_query_cache[0][0, :, 0, 0], + torch.tensor([4, 5], dtype=torch.bfloat16), + ) + + manager.free_part_slots( + 0, + owner, + torch.tensor([0, 2, 4, 5], dtype=torch.long), + keep_indices_sorted=True, + ) + assert manager._rkv_query_positions[0][0].tolist() == [-1, -1] + resident_slots = manager.buffer_req_to_token_slots[0][0, :4].clone().long() + payload = storage.layer_payload(0) + torch.testing.assert_close( + payload.latent_cache[resident_slots, 0, 0], + torch.tensor([0, 2, 4, 5], dtype=torch.bfloat16), + ) + torch.testing.assert_close( + payload.rope_cache[resident_slots, 0, 0], + torch.tensor([100, 102, 104, 105], dtype=torch.bfloat16), + ) + coordinator.index.finish( + owner.chain_id, + token_ids=owner_tokens, + processed_token_count=len(owner_tokens), + physical_slots_by_layer=manager.chain_physical_residency(owner.seq_id), + ) + + resumed_tokens = owner_tokens + [6, 7] + resumed_plan = coordinator.plan_admission( + chain_id=owner.chain_id, + seq_id=owner.seq_id, + token_ids=resumed_tokens, + ) + assert resumed_plan.status == "resumed" + assert resumed_plan.reused_tokens == len(owner_tokens) + coordinator.apply_admission(resumed_plan) + resumed = Sequence(resumed_tokens) + resumed.seq_id = owner.seq_id + resumed.chain_id = owner.chain_id + resumed.chain_status = "resumed" + resumed.chain_reused_tokens = len(owner_tokens) + resumed.num_prefilled_tokens = len(owner_tokens) + resumed.current_chunk_size = 2 + + input_ids, positions, _ = manager._prepare_prefill([resumed]) + assert input_ids.tolist() == [6, 7] + assert positions.tolist() == [6, 7] + assert manager._rkv_query_positions[0][0].tolist() == [-1, -1] + resumed_row = manager.seq_id_to_row[0][resumed.seq_id] + resumed_slots = manager.buffer_req_to_token_slots[0][ + resumed_row, :6 + ].clone().long() + assert resumed_slots[:4].tolist() == resident_slots.tolist() + _fill_latent_slots(storage, resumed_slots[4:], [6, 7]) + resumed_q = torch.tensor([106, 107], dtype=torch.bfloat16).view( + 2, 1, 1 + ).expand(2, 2, 256) + manager.record_prefill_query( + 0, + resumed_q, + _rkv_prefill_view(manager, storage), + b_start_loc=torch.tensor([0], dtype=torch.int32), + chunk_lens=torch.tensor([2], dtype=torch.int32), + ) + assert manager._rkv_query_positions[0][0].tolist() == [4, 5] + torch.testing.assert_close( + manager._rkv_query_cache[0][0, :, 0, 0], + torch.tensor([106, 107], dtype=torch.bfloat16), + ) + manager.free_part_slots( + 0, + resumed, + torch.tensor([0, 2, 4, 5], dtype=torch.long), + keep_indices_sorted=True, + ) + assert manager._rkv_query_positions[0][0].tolist() == [-1, -1] + final_slots = manager.buffer_req_to_token_slots[0][ + resumed_row, :4 + ].clone().long() + torch.testing.assert_close( + payload.latent_cache[final_slots, 0, 0], + torch.tensor([0, 4, 6, 7], dtype=torch.bfloat16), + ) + torch.testing.assert_close( + payload.rope_cache[final_slots, 0, 0], + torch.tensor([100, 104, 106, 107], dtype=torch.bfloat16), + ) + + coordinator.index.finish( + owner.chain_id, + token_ids=resumed_tokens, + processed_token_count=len(resumed_tokens), + physical_slots_by_layer=manager.chain_physical_residency(resumed.seq_id), + ) + coordinator.invalidate(owner.chain_id) + manager.free_seq(resumed.seq_id) + assert manager._rkv_query_positions[0][0].tolist() == [-1, -1] + assert manager._prefill_attn_score_accumulators == {} + assert manager.seq_id_to_row == [{}] + assert manager._num_free_slots == [32] diff --git a/tests/test_glm_mla_sparse_methods.py b/tests/test_glm_mla_sparse_methods.py new file mode 100644 index 00000000..1e8f7f01 --- /dev/null +++ b/tests/test_glm_mla_sparse_methods.py @@ -0,0 +1,343 @@ +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import patch + +import numpy as np +import pytest +import torch +from transformers import Glm4MoeLiteConfig + +from sparsevllm.config import RuntimeLayout +from sparsevllm.engine.cache_manager import LayerBatchStates +from sparsevllm.engine.cache_manager.snapkv import SnapKVCacheManager +from sparsevllm.engine.cache_manager.rkv import RKVCacheManager +from sparsevllm.engine.cache_manager.standard import StandardCacheManager +from sparsevllm.engine.cache_manager.storage import MlaLatentStorage +from sparsevllm.engine.cache_manager.streamingllm import ( + StreamingLLMCacheManager, +) +from sparsevllm.engine.sequence import Sequence +from sparsevllm.engine.sparse_controller import ( + LayerBatchSparseState, + SparseController, +) +from sparsevllm.utils.context import reset_context, set_context + +from glm_test_helpers import _single_rank_parallel_context + + +def test_glm_sparse_controller_uses_full_mla_qk_softmax_scale(): + hf_config = Glm4MoeLiteConfig( + hidden_size=64, + num_attention_heads=20, + num_key_value_heads=20, + qk_nope_head_dim=192, + qk_rope_head_dim=64, + ) + config = SimpleNamespace( + vllm_sparse_method="omnikv", + obs_layer_ids=[], + full_attn_layers=[], + runtime_layout=RuntimeLayout.dense(1), + hf_config=hf_config, + tensor_parallel_size=1, + num_sink_tokens=0, + num_recent_tokens=0, + decode_keep_tokens=1, + sparse_attn_score_dtype="float32", + ) + manager = SimpleNamespace(device=torch.device("cpu")) + + controller = SparseController(config, manager) + + assert hf_config.head_dim == 64 + assert controller.attn_softmax_scale == pytest.approx(256**-0.5) + assert controller.attn_softmax_scale != pytest.approx(64**-0.5) + + +def test_glm_rkv_query_cache_allocates_and_records_full_qk_head_width(): + hf_config = Glm4MoeLiteConfig( + hidden_size=64, + num_hidden_layers=1, + num_attention_heads=20, + num_key_value_heads=20, + qk_nope_head_dim=192, + qk_rope_head_dim=64, + torch_dtype=torch.bfloat16, + ) + config = SimpleNamespace( + hf_config=hf_config, + runtime_layout=RuntimeLayout.dense(1), + vllm_sparse_method="rkv", + max_model_len=8, + max_num_seqs_in_gpu=1, + num_kvcache_slots=8, + pyramid_layer_ratios=None, + num_sink_tokens=1, + num_recent_tokens=1, + decode_keep_tokens=1, + rkv_compression_interval=1, + rkv_observation_tokens=1, + ) + cpu_platform = SimpleNamespace( + get_device=lambda _rank: torch.device("cpu"), + supports_pin_memory=lambda: False, + ) + with ( + patch( + "sparsevllm.engine.cache_manager.base.platforms.get_current_platform", + return_value=cpu_platform, + ), + patch( + "sparsevllm.engine.cache_manager.snapkv.create_attention_cache_storage", + return_value=SimpleNamespace(), + ), + patch.object(SnapKVCacheManager, "allocate_kv_cache", autospec=True), + ): + manager = RKVCacheManager(config, _single_rank_parallel_context()) + + assert hf_config.head_dim == 64 + assert manager.head_dim == 256 + assert manager._rkv_query_cache[0].shape == (1, 1, 20, 256) + manager.layer_batch_states[0] = LayerBatchStates( + req_indices=torch.tensor([0], dtype=torch.int32), + context_lens=torch.tensor([4], dtype=torch.int32), + ) + q = torch.arange(20 * 256, dtype=torch.float32).reshape(1, 20, 256) + q = q.to(torch.bfloat16) + + manager.record_decode_query(0, q) + + assert manager._rkv_query_positions[0].tolist() == [[3]] + torch.testing.assert_close(manager._rkv_query_cache[0][0, 0], q[0]) + + +def _latent_snap_family_manager(manager_type, *, row_len: int): + storage = MlaLatentStorage( + kv_lora_rank=512, + rope_dim=64, + dtype=torch.bfloat16, + ) + storage.allocate(num_layers=1, num_slots=16, device=torch.device("cpu")) + manager = object.__new__(manager_type) + manager.attention_cache_storage = storage + manager.kv_cache = None + manager.device = torch.device("cpu") + manager.num_layers = 1 + manager.num_kv_layers = 1 + manager.runtime_layout = RuntimeLayout.dense(1) + manager._uniform_decode_metadata = True + manager.buffer_req_to_token_slots_tensor = torch.zeros( + (1, 1, 16), dtype=torch.int32 + ) + manager.buffer_req_to_token_slots_tensor[0, 0, :row_len] = torch.arange( + row_len, dtype=torch.int32 + ) + manager.buffer_req_to_token_slots = [ + manager.buffer_req_to_token_slots_tensor[0] + ] + manager.seq_id_to_row = [{0: 0}] + manager.row_seq_lens = [np.asarray([row_len], dtype=np.int32)] + manager.free_slots_stack_tensor = None + manager.free_slots_stack = [torch.zeros((16,), dtype=torch.int32)] + manager._num_free_slots = [0] + + payload = storage.layer_payload(0) + for slot in range(row_len): + payload.latent_cache[slot].fill_(slot) + payload.rope_cache[slot].fill_(slot + 100) + return manager, payload + + +def test_streamingllm_budget_trigger_preserves_mla_latent_slot_payloads(): + manager, payload = _latent_snap_family_manager( + StreamingLLMCacheManager, + row_len=8, + ) + + seq = Sequence(list(range(8))) + seq.seq_id = 0 + seq.num_prefilled_tokens = 0 + seq.current_chunk_size = 8 + controller = object.__new__(SparseController) + controller.cache_manager = manager + controller.device = torch.device("cpu") + controller.num_layers = 1 + controller.num_sink = 2 + controller.num_recent = 3 + controller.layer_batch_sparse_states = { + 0: SimpleNamespace( + context_lens=torch.tensor([8], dtype=torch.int32), + max_context_len=8, + ) + } + controller._is_kv_layer = lambda layer_idx: int(layer_idx) == 0 + + controller._streamingllm_prefill_eviction([seq]) + + active_slots = manager.buffer_req_to_token_slots[0][0, :5].long() + assert active_slots.tolist() == [0, 1, 5, 6, 7] + assert manager.row_seq_lens[0].tolist() == [5] + assert manager.free_slots_stack[0][:3].tolist() == [2, 3, 4] + assert manager._num_free_slots == [3] + torch.testing.assert_close( + payload.latent_cache[active_slots, 0, 0], + torch.tensor([0, 1, 5, 6, 7], dtype=torch.bfloat16), + ) + torch.testing.assert_close( + payload.rope_cache[active_slots, 0, 0], + torch.tensor([100, 101, 105, 106, 107], dtype=torch.bfloat16), + ) + + +def test_snapkv_score_budget_trigger_preserves_mla_latent_slot_payloads(): + manager, payload = _latent_snap_family_manager( + SnapKVCacheManager, + row_len=8, + ) + manager._prefill_attn_score_accumulators = { + (0, 0): torch.tensor([0.0, 1.0, 2.0, 9.0, 3.0, 8.0, 4.0, 0.0]) + } + seq = Sequence(list(range(8))) + seq.seq_id = 0 + seq.num_prefilled_tokens = 0 + seq.current_chunk_size = 8 + controller = object.__new__(SparseController) + controller.cache_manager = manager + controller.device = torch.device("cpu") + controller.sparse_method = "snapkv" + controller.num_layers = 1 + controller.num_sink = 1 + controller.num_recent = 1 + controller.decode_keep_tokens = 2 + controller.config = SimpleNamespace( + snapkv_num_full_layers=0, + pyramid_layer_ratios=None, + pool_kernel_size=1, + ) + controller._is_kv_layer = lambda layer_idx: int(layer_idx) == 0 + controller._kv_layer_index = lambda layer_idx: int(layer_idx) + + controller._snapkv_prefill_eviction([seq]) + + active_slots = manager.buffer_req_to_token_slots[0][0, :4].long() + assert active_slots.tolist() == [0, 3, 5, 7] + assert manager.row_seq_lens[0].tolist() == [4] + assert sorted(manager.free_slots_stack[0][:4].tolist()) == [1, 2, 4, 6] + assert manager._num_free_slots == [4] + assert manager._prefill_attn_score_accumulators == {} + torch.testing.assert_close( + payload.latent_cache[active_slots, 0, 0], + torch.tensor([0, 3, 5, 7], dtype=torch.bfloat16), + ) + torch.testing.assert_close( + payload.rope_cache[active_slots, 0, 0], + torch.tensor([100, 103, 105, 107], dtype=torch.bfloat16), + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_omnikv_observation_selects_mla_latent_active_slots(): + device = torch.device("cuda") + storage = MlaLatentStorage( + kv_lora_rank=512, + rope_dim=64, + dtype=torch.bfloat16, + ) + storage.allocate(num_layers=2, num_slots=16, device=device) + manager = object.__new__(StandardCacheManager) + manager.attention_cache_storage = storage + manager.kv_cache = None + manager.runtime_layout = RuntimeLayout.dense(2) + manager.buffer_req_to_token_slots = torch.zeros( + (1, 16), dtype=torch.int32, device=device + ) + physical_slots = torch.tensor( + [8, 3, 11, 1, 14, 7], dtype=torch.int32, device=device + ) + manager.buffer_req_to_token_slots[0, :6] = physical_slots + payload = storage.layer_payload(1) + for slot in physical_slots.tolist(): + payload.latent_cache[slot].fill_(slot) + payload.rope_cache[slot].fill_(slot + 100) + + hf_config = Glm4MoeLiteConfig( + hidden_size=64, + num_hidden_layers=2, + num_attention_heads=20, + num_key_value_heads=20, + qk_nope_head_dim=192, + qk_rope_head_dim=64, + torch_dtype=torch.bfloat16, + ) + controller = SparseController( + SimpleNamespace( + vllm_sparse_method="omnikv", + obs_layer_ids=[0], + full_attn_layers=[0], + runtime_layout=RuntimeLayout.dense(2), + hf_config=hf_config, + tensor_parallel_size=1, + num_sink_tokens=1, + num_recent_tokens=1, + decode_keep_tokens=2, + sparse_attn_score_dtype="float32", + ), + manager, + ) + assert controller.attn_softmax_scale == pytest.approx(256**-0.5) + controller.layer_batch_sparse_states = { + 0: LayerBatchSparseState( + attn_score=torch.tensor( + [ + [ + [0.0, 10.0, 1.0, 9.0, 0.0, -1.0], + [0.0, 0.0, 8.0, 1.0, 7.0, -1.0], + ] + ], + dtype=torch.float32, + device=device, + ), + req_indices=torch.tensor([0], dtype=torch.int32, device=device), + context_lens=torch.tensor([6], dtype=torch.int32, device=device), + max_context_len=6, + ), + 1: LayerBatchSparseState(), + } + controller._is_kv_layer = lambda layer_idx: 0 <= int(layer_idx) < 2 + + set_context( + False, + cache_manager=manager, + is_long_text=True, + seqs=[Sequence([1])], + ) + try: + controller.on_layer_end(0, SimpleNamespace(is_prefill=False)) + finally: + reset_context() + + target = controller.layer_batch_sparse_states[1] + assert target.active_indices is not None + assert target.active_slots is not None + assert target.context_lens is not None + assert target.context_lens.tolist() == [4] + logical_keep = target.active_indices[0, :4].tolist() + assert logical_keep[0] == 0 + assert set(logical_keep[1:3]) == {1, 2} + assert logical_keep[3] == 5 + selected_slots = target.active_slots[0, :4].long() + expected_slots = manager.buffer_req_to_token_slots[0].index_select( + 0, + target.active_indices[0, :4].long(), + ) + torch.testing.assert_close(selected_slots, expected_slots.long()) + torch.testing.assert_close( + payload.latent_cache[selected_slots, 0, 0], + selected_slots.to(torch.bfloat16), + ) + torch.testing.assert_close( + payload.rope_cache[selected_slots, 0, 0], + selected_slots.to(torch.bfloat16) + 100, + ) diff --git a/tests/test_glm_runtime_compatibility.py b/tests/test_glm_runtime_compatibility.py new file mode 100644 index 00000000..d11a4f82 --- /dev/null +++ b/tests/test_glm_runtime_compatibility.py @@ -0,0 +1,297 @@ +from __future__ import annotations + +import pytest +import torch + +from sparsevllm.engine.cache_manager.storage import CacheLayout +from sparsevllm.method_registry import ( + GLM4_MOE_LITE_EP_COMPATIBILITY, + MODEL_RUNTIME_COMPATIBILITY, +) +from sparsevllm.distributed import ParallelMode + +from glm_test_helpers import _glm_config + + +def test_glm_config_selects_mla_latent_layout(): + config = _glm_config() + + assert config.attention_cache_layout == CacheLayout.MLA_LATENT.value + assert config.mla_prefill_workspace_bytes == 2 * 1024**3 + + +def test_mla_prefill_workspace_budget_must_be_positive(): + with pytest.raises(ValueError, match="mla_prefill_workspace_bytes"): + _glm_config(mla_prefill_workspace_bytes=0) + + +@pytest.mark.parametrize( + ("override", "error_type", "message"), + [ + ({"vllm_sparse_method": "quest"}, ValueError, "Unsupported glm4_moe_lite"), + ( + {"enforce_eager": False}, + ValueError, + "requires enforce_eager=True", + ), + ], +) +def test_glm_config_rejects_unsupported_storage_combinations( + override, error_type, message +): + with pytest.raises(error_type, match=message): + _glm_config(**override) + + +@pytest.mark.parametrize("expert_parallel_size", [2, 4]) +def test_glm_config_accepts_replicated_attention_ep(expert_parallel_size): + config = _glm_config( + tensor_parallel_size=1, + expert_parallel_size=expert_parallel_size, + data_parallel_size=1, + ) + + assert config.world_size == expert_parallel_size + assert config.tensor_parallel_size == 1 + assert config.expert_parallel_size == expert_parallel_size + assert config.data_parallel_size == 1 + assert not config.uses_outer_tp_moe_layout + + +@pytest.mark.parametrize( + ( + "tensor_parallel_size", + "expert_parallel_size", + "world_size", + "moe_tensor_parallel_size", + ), + [ + (2, 2, 2, 1), + (4, 2, 4, 2), + (4, 4, 4, 1), + ], +) +def test_glm_config_accepts_outer_tp_moe_ep_layout( + tensor_parallel_size, + expert_parallel_size, + world_size, + moe_tensor_parallel_size, +): + config = _glm_config( + tensor_parallel_size=tensor_parallel_size, + expert_parallel_size=expert_parallel_size, + ) + + assert config.uses_outer_tp_moe_layout + assert config.world_size == world_size + assert config.moe_tensor_parallel_size == moe_tensor_parallel_size + + +def test_glm_hybrid_checks_routed_width_against_moe_tp_not_outer_tp(): + config = _glm_config( + tensor_parallel_size=4, + expert_parallel_size=2, + hf_overrides={"moe_intermediate_size": 6}, + ) + + assert config.moe_tensor_parallel_size == 2 + + +def test_glm_hybrid_rejects_routed_width_not_divisible_by_moe_tp(): + with pytest.raises(ValueError, match="divisible by MoE TP"): + _glm_config( + tensor_parallel_size=4, + expert_parallel_size=2, + hf_overrides={"moe_intermediate_size": 5}, + ) + + +def test_glm_config_rejects_nondivisible_outer_tp_moe_ep_layout(): + with pytest.raises(ValueError, match="TP divisible by EP"): + _glm_config(tensor_parallel_size=2, expert_parallel_size=4) + + +def test_glm_config_rejects_data_parallelism(): + with pytest.raises(ValueError, match="does not support data parallelism"): + _glm_config(data_parallel_size=2) + + +_GLM_PARALLEL_LAYOUTS = [ + (1, 1), + (2, 1), + (4, 1), + (1, 2), + (1, 4), + (2, 2), + (4, 2), + (4, 4), +] + + +@pytest.mark.parametrize( + ("tensor_parallel_size", "expert_parallel_size"), + _GLM_PARALLEL_LAYOUTS, +) +@pytest.mark.parametrize( + "method", + ["", "streamingllm", "snapkv", "h2o", "omnikv", "rkv"], +) +def test_glm_config_accepts_parallel_sparse_graph_cross_product( + tensor_parallel_size, + expert_parallel_size, + method, +): + config = _glm_config( + tensor_parallel_size=tensor_parallel_size, + expert_parallel_size=expert_parallel_size, + vllm_sparse_method=method, + decode_cuda_graph=True, + ) + + assert config.decode_cuda_graph + assert config.vllm_sparse_method == method + + +@pytest.mark.parametrize( + ("tensor_parallel_size", "expert_parallel_size"), + _GLM_PARALLEL_LAYOUTS, +) +@pytest.mark.parametrize( + "method", + ["", "streamingllm", "snapkv", "h2o", "omnikv", "rkv"], +) +def test_glm_config_accepts_parallel_sparse_graph_prefix_cross_product( + tensor_parallel_size, + expert_parallel_size, + method, +): + config = _glm_config( + tensor_parallel_size=tensor_parallel_size, + expert_parallel_size=expert_parallel_size, + vllm_sparse_method=method, + decode_cuda_graph=True, + enable_prefix_caching=True, + ) + + assert config.enable_prefix_caching + assert config.decode_cuda_graph + assert config.resolved_prefix_cache_mode == ( + "radix" if method in {"", "omnikv"} else "chain" + ) + + +def test_glm_registry_selects_parallel_layout_contracts(): + assert MODEL_RUNTIME_COMPATIBILITY[ + ("glm4_moe_lite", ParallelMode.STANDARD) + ] is GLM4_MOE_LITE_EP_COMPATIBILITY + assert MODEL_RUNTIME_COMPATIBILITY[ + ("glm4_moe_lite", ParallelMode.OUTER_TP_MOE) + ] is GLM4_MOE_LITE_EP_COMPATIBILITY + + +@pytest.mark.parametrize( + "method", + ["", "streamingllm", "snapkv", "h2o", "omnikv", "rkv"], +) +def test_glm_config_accepts_tp1_ep1_decode_cuda_graph(method): + config = _glm_config( + decode_cuda_graph=True, + vllm_sparse_method=method, + ) + + assert config.decode_cuda_graph + assert config.tensor_parallel_size == 1 + assert config.expert_parallel_size == 1 + assert config.vllm_sparse_method == method + + +def test_glm_config_accepts_vanilla_latent_prefix_cache(): + config = _glm_config(enable_prefix_caching=True) + + assert config.attention_cache_layout == CacheLayout.MLA_LATENT.value + assert config.enable_prefix_caching + assert config.resolved_prefix_cache_mode == "radix" + + +def test_glm_config_accepts_omnikv_latent_prefix_cache_after_lifecycle_gate(): + config = _glm_config( + vllm_sparse_method="omnikv", + enable_prefix_caching=True, + ) + + assert config.attention_cache_layout == CacheLayout.MLA_LATENT.value + assert config.enable_prefix_caching + assert config.resolved_prefix_cache_mode == "radix" + + +def test_glm_config_accepts_snapkv_latent_chain_prefix_after_lifecycle_gate(): + config = _glm_config( + vllm_sparse_method="snapkv", + enable_prefix_caching=True, + ) + + assert config.attention_cache_layout == CacheLayout.MLA_LATENT.value + assert config.enable_prefix_caching + assert config.resolved_prefix_cache_mode == "chain" + + +def test_glm_config_accepts_h2o_latent_chain_prefix_after_lifecycle_gate(): + config = _glm_config( + vllm_sparse_method="h2o", + enable_prefix_caching=True, + ) + + assert config.attention_cache_layout == CacheLayout.MLA_LATENT.value + assert config.enable_prefix_caching + assert config.resolved_prefix_cache_mode == "chain" + + +def test_glm_config_accepts_rkv_latent_chain_prefix_after_lifecycle_gate(): + config = _glm_config( + vllm_sparse_method="rkv", + enable_prefix_caching=True, + ) + + assert config.attention_cache_layout == CacheLayout.MLA_LATENT.value + assert config.enable_prefix_caching + assert config.resolved_prefix_cache_mode == "chain" + + +def test_glm_config_accepts_streamingllm_chain_prefix_cache(): + config = _glm_config( + vllm_sparse_method="streamingllm", + enable_prefix_caching=True, + ) + + assert config.resolved_prefix_cache_mode == "chain" + + +@pytest.mark.parametrize( + "method", + ["streamingllm", "snapkv", "h2o", "omnikv", "rkv"], +) +def test_glm_config_accepts_sparse_latent_layout(method): + config = _glm_config( + vllm_sparse_method=method, + num_sink_tokens=2, + num_recent_tokens=3, + ) + + assert config.attention_cache_layout == CacheLayout.MLA_LATENT.value + assert config.vllm_sparse_method == method + + +@pytest.mark.parametrize( + ("hf_override", "message"), + [ + ({"torch_dtype": torch.float16}, "requires BF16"), + ({"kv_lora_rank": 256}, "requires kv_lora_rank=512"), + ({"qk_rope_head_dim": 32}, "qk_rope_head_dim=64"), + ], +) +def test_glm_config_rejects_unsupported_mla_storage_contract( + hf_override, + message, +): + with pytest.raises(NotImplementedError, match=message): + _glm_config(hf_overrides=hf_override) diff --git a/tests/test_h2o_cache_manager.py b/tests/test_h2o_cache_manager.py index 2a7530b9..a182c650 100644 --- a/tests/test_h2o_cache_manager.py +++ b/tests/test_h2o_cache_manager.py @@ -17,6 +17,7 @@ ) from sparsevllm.engine.cache_manager.h2o import H2OCacheManager from sparsevllm.engine.cache_manager.snapkv import SnapKVCacheManager +from sparsevllm.engine.cache_manager.storage import MlaLatentStorage from sparsevllm.engine.scheduler import Scheduler from sparsevllm.engine.sequence import Sequence from sparsevllm.engine.sparse_controller import SparseController @@ -802,6 +803,45 @@ def tracked_workspace(**kwargs): assert workspace_entries[0]["nbytes"] == workspace.untyped_storage().nbytes() +def test_h2o_final_prefill_compacts_mla_latent_and_rope_slots(): + manager = _manager_with_rows([6], decode_budget=4, prefill_budget=8) + _set_layer_row_slots(manager, 0, [[9, 2, 7, 1, 6, 4]]) + manager._h2o_scores[(0, 0)] = torch.tensor( + [1.0, 9.0, 2.0, 8.0, 0.0, 0.0] + ) + storage = MlaLatentStorage( + kv_lora_rank=512, + rope_dim=64, + dtype=torch.bfloat16, + ) + storage.allocate(num_layers=1, num_slots=64, device=torch.device("cpu")) + manager.attention_cache_storage = storage + manager.kv_cache = None + assert storage.latent_cache is not None + assert storage.rope_cache is not None + for slot in [9, 2, 7, 1, 6, 4]: + storage.latent_cache[0, slot].fill_(slot) + storage.rope_cache[0, slot].fill_(slot + 100) + seq = _seq(0, 6, prefilled=0, chunk=6) + + manager.evict_after_prefill([seq]) + + destination_slots = manager.buffer_req_to_token_slots[0][0, :4].long() + assert destination_slots.tolist() == [1, 2, 4, 6] + payload = storage.layer_payload(0) + expected_sources = torch.tensor([2, 1, 6, 4], dtype=torch.bfloat16) + torch.testing.assert_close( + payload.latent_cache[destination_slots, 0, 0], + expected_sources, + ) + torch.testing.assert_close( + payload.rope_cache[destination_slots, 0, 0], + expected_sources + 100, + ) + assert manager._h2o_scores[(0, 0)].tolist() == [9.0, 8.0, 0.0, 0.0] + assert manager._h2o_final_prefill_workspace is None + + def test_h2o_intermediate_prefill_does_not_move_kv_payloads(): manager = _manager_with_rows([6], decode_budget=3, prefill_budget=4) _set_layer_row_slots(manager, 0, [[9, 2, 7, 1, 6, 4]]) diff --git a/tests/test_longbench_deltakv_contracts.py b/tests/test_longbench_deltakv_contracts.py index 1941209d..ea88f101 100644 --- a/tests/test_longbench_deltakv_contracts.py +++ b/tests/test_longbench_deltakv_contracts.py @@ -10,6 +10,22 @@ class LongBenchDeltaKVContractsTest(unittest.TestCase): + def test_longbench_accepts_explicit_runtime_context_limit(self): + with patch.object( + longbench_pred.sys, + "argv", + [ + "pred.py", + "--model_path", + "/models/glm", + "--max_model_len", + "32768", + ], + ): + args = longbench_pred.parse_args() + + self.assertEqual(args.max_model_len, 32768) + def test_no_chat_datasets_remain_raw_for_every_thinking_mode(self): for dataset in longbench_pred.NO_CHAT_TEMPLATE_DATASETS: for thinking_mode in ("off", "on", "on_strip"): @@ -101,6 +117,10 @@ def test_longbench_records_actual_decode_cuda_graph_state(self): "uncaptured": SimpleNamespace(graph=None), }, last_state_key="captured", + capture_count=2, + replay_count=7, + eager_static_count=0, + force_eager_count=0, ) generate_fn = SimpleNamespace( _sparsevllm_llm=SimpleNamespace( @@ -125,12 +145,51 @@ def test_longbench_records_actual_decode_cuda_graph_state(self): self.assertEqual(status["state_count"], 2) self.assertEqual(status["graph_count"], 1) self.assertTrue(status["active"]) + self.assertEqual(status["capture_count"], 2) + self.assertEqual(status["replay_count"], 7) self.assertEqual(status["last_state_key"], "captured") self.assertEqual( json.loads(path.read_text(encoding="utf-8")), status, ) + def test_longbench_records_business_graph_counter_delta(self): + graph_runner = SimpleNamespace( + _graphs={"captured": SimpleNamespace(graph=object())}, + last_state_key="captured", + capture_count=1, + replay_count=3, + eager_static_count=0, + force_eager_count=0, + ) + generate_fn = SimpleNamespace( + _sparsevllm_llm=SimpleNamespace( + config=SimpleNamespace(decode_cuda_graph=True), + model_runner=SimpleNamespace( + decode_cuda_graph_runner=graph_runner, + ), + ) + ) + before = longbench_pred._decode_cuda_graph_status( + generate_fn=generate_fn, + rank=0, + ) + graph_runner.replay_count = 11 + + with tempfile.TemporaryDirectory() as tmp: + status = longbench_pred._write_decode_cuda_graph_status( + generate_fn=generate_fn, + out_root=tmp, + rank=0, + before=before, + ) + + self.assertEqual(status["before"]["replay_count"], 3) + self.assertEqual(status["replay_count"], 11) + self.assertEqual(status["counter_delta"]["replay_count"], 8) + self.assertEqual(status["counter_delta"]["eager_static_count"], 0) + self.assertEqual(status["counter_delta"]["force_eager_count"], 0) + def test_longbench_fails_if_sparsevllm_graph_state_is_unavailable(self): with tempfile.TemporaryDirectory() as tmp: with self.assertRaisesRegex(RuntimeError, "_sparsevllm_llm"): diff --git a/tests/test_mla_attention_layer.py b/tests/test_mla_attention_layer.py index 779245b7..2647738b 100644 --- a/tests/test_mla_attention_layer.py +++ b/tests/test_mla_attention_layer.py @@ -7,6 +7,7 @@ import sparsevllm.layers.mla_attention as mla_attention_module from sparsevllm.engine.cache_manager import ( + AttentionKeyComputeView, AttentionViewMeta, ExplicitKVPayload, MlaLatentPayload, @@ -20,7 +21,7 @@ MlaAttentionOpSpec, MlaAttentionProvider, ) -from sparsevllm.utils.context import reset_context +from sparsevllm.utils.context import get_context, reset_context, set_context class _TestProvider(MlaAttentionProvider): @@ -240,6 +241,18 @@ def test_mla_attention_bind_resolves_provider_once() -> None: ) +def test_set_context_starts_a_new_attention_validation_scope() -> None: + reset_context() + initial_scope = get_context().attention_validation_scope + set_context(False) + first_step_scope = get_context().attention_validation_scope + set_context(False) + second_step_scope = get_context().attention_validation_scope + + assert first_step_scope is not initial_scope + assert second_step_scope is not first_step_scope + + def test_mla_prefill_reuses_validated_packing_across_layers() -> None: attention = _attention() active_slots = torch.tensor([[3, 1]], dtype=torch.int32) @@ -330,6 +343,84 @@ def test_mla_prefill_reuses_query_validation_across_layers() -> None: assert host_values.call_count == 4 +def test_mla_prefill_exposes_expanded_explicit_score_view() -> None: + attention = _attention() + view = _view( + torch.empty(2, 1, 512, dtype=torch.bfloat16), + torch.empty(2, 1, 64, dtype=torch.bfloat16), + torch.tensor([[0, 1]], dtype=torch.int32), + torch.tensor([0], dtype=torch.int32), + torch.tensor([2], dtype=torch.int32), + ) + with patch("sparsevllm.layers.mla_attention.gather_latent_history"): + history = attention.prepare_prefill_history(view, query_tokens=2) + workset = _expand_history(attention, history) + + score_view = attention.build_prefill_explicit_view(workset) + + assert isinstance(score_view.payload, ExplicitKVPayload) + assert score_view.payload.k_cache is workset.expanded_k + assert score_view.payload.v_cache is workset.expanded_v + assert score_view.meta.active_slots is history.packed_slots + assert score_view.meta.req_indices is history.local_req_indices + assert score_view.meta.context_lens is history.context_lens + assert score_view.meta.max_context_len == history.max_context_len + + +def test_mla_materializes_actual_keys_for_permuted_slots() -> None: + torch.manual_seed(47) + attention = _attention(tp_size=4) + latent_cache = torch.randn(7, 1, 512, dtype=torch.bfloat16) + rope_cache = torch.randn(7, 1, 64, dtype=torch.bfloat16) + slots = torch.tensor([[5, 1], [6, 2]], dtype=torch.int32) + view = AttentionKeyComputeView( + active_slots=slots, + payload=MlaLatentPayload( + latent_cache=latent_cache, + rope_cache=rope_cache, + ), + ) + weights = torch.randn( + attention.spec.local_q_heads, + 448, + 512, + dtype=torch.bfloat16, + ) + + def project_latent(latent: torch.Tensor) -> torch.Tensor: + return torch.einsum( + "tr,hor->tho", + latent.float(), + weights.float(), + ).to(torch.bfloat16).flatten(1) + + actual = attention.materialize_expanded_keys( + view, + project_latent=project_latent, + ) + + flat_slots = slots.long().flatten() + latent = latent_cache[flat_slots, 0] + projected = torch.einsum( + "tr,hor->tho", + latent.float(), + weights.float(), + ).to(torch.bfloat16) + expected = torch.cat( + ( + projected[..., :192], + rope_cache[flat_slots, 0][:, None, :].expand( + -1, + attention.spec.local_q_heads, + -1, + ), + ), + dim=-1, + ).view(2, 2, attention.spec.local_q_heads, 256) + + torch.testing.assert_close(actual, expected) + + CUDA_REQUIRED = pytest.mark.skipif( not torch.cuda.is_available(), reason="CUDA is required for MLA full-history prefill tests", diff --git a/tests/test_mla_attention_operator.py b/tests/test_mla_attention_operator.py index 3bf28897..c27f32ab 100644 --- a/tests/test_mla_attention_operator.py +++ b/tests/test_mla_attention_operator.py @@ -115,6 +115,26 @@ def test_mla_resolver_selects_h100_triton_provider() -> None: ) +def test_mla_resolver_accepts_cuda_graph_after_capture_gate() -> None: + spec = _spec(cuda_graph=True) + workspace = _cpu_workspace(batch_size=8, head_count=spec.local_q_heads) + + with patch( + "sparsevllm.operators.mla_attention.allocate_mla_decode_workspace", + return_value=workspace, + ): + resolved = OpResolver(MLA_ATTENTION_REGISTRY).resolve( + spec, + _h100_caps(), + op_spec=spec, + device="cpu", + max_batch_size=8, + ) + + assert isinstance(resolved.provider, MlaTritonProvider) + assert resolved.rejected == () + + @pytest.mark.parametrize( ("spec_overrides", "caps_overrides", "reason"), [ @@ -123,11 +143,15 @@ def test_mla_resolver_selects_h100_triton_provider() -> None: ({}, {"device_name": "NVIDIA H100 PCIe"}, "validated"), ({}, {"supports_triton": False}, "does not support Triton"), ({}, {"supports_bfloat16": False}, "does not support BF16"), + ( + {"cuda_graph": True}, + {"supports_graph_capture": False}, + "graph capture support", + ), ({"activation_dtype": torch.float16}, {}, "BF16 activations"), ({"cache_dtype": torch.float16}, {}, "BF16 cache"), ({"kv_lora_rank": 256}, {}, "GLM MLA shape"), ({"tp_size": 5}, {}, "tensor parallel size"), - ({"cuda_graph": True}, {}, "CUDA Graph"), ], ) def test_mla_resolver_rejects_unvalidated_contracts( @@ -252,6 +276,7 @@ def test_mla_provider_run_does_not_resolve_or_allocate() -> None: view.meta.req_indices, view.meta.context_lens, cache_slot_count=4, + max_context_len=None, ) assert kernel.call_count == 3 kernel.assert_called_with( @@ -265,11 +290,69 @@ def test_mla_provider_run_does_not_resolve_or_allocate() -> None: output, workspace, softmax_scale=spec.softmax_scale, + attn_score=None, + max_context_len=None, config=provider.launch_config, validate_metadata=False, ) +def test_mla_provider_validates_each_metadata_identity_once_per_scope() -> None: + spec = _spec(tp_size=4) + workspace = _cpu_workspace(batch_size=1, head_count=5) + with patch( + "sparsevllm.operators.mla_attention.allocate_mla_decode_workspace", + return_value=workspace, + ): + provider = MlaTritonProvider( + op_spec=spec, + device="cpu", + max_batch_size=1, + ) + + payload = MlaLatentPayload( + latent_cache=torch.empty(4, 1, 512, dtype=torch.bfloat16), + rope_cache=torch.empty(4, 1, 64, dtype=torch.bfloat16), + ) + + def view(slots: list[int]) -> DecodeComputeView: + return DecodeComputeView( + meta=AttentionViewMeta( + active_slots=torch.tensor([slots], dtype=torch.int32), + req_indices=torch.tensor([0], dtype=torch.int32), + context_lens=torch.tensor([len(slots)], dtype=torch.int32), + ), + payload=payload, + ) + + view_a = view([0, 1]) + view_b = view([2, 3]) + q_nope_absorbed = torch.empty(1, 5, 512, dtype=torch.bfloat16) + q_rope = torch.empty(1, 5, 64, dtype=torch.bfloat16) + output = torch.empty_like(q_nope_absorbed) + validation_scope = object() + + with ( + patch( + "sparsevllm.operators.mla_attention.run_mla_decode", + return_value=output, + ), + patch( + "sparsevllm.operators.mla_attention.validate_mla_decode_metadata" + ) as validate, + ): + for decode_view in (view_a, view_b, view_a, view_b): + provider.run( + q_nope_absorbed, + q_rope, + decode_view, + output, + validation_scope=validation_scope, + ) + + assert validate.call_count == 2 + + def test_mla_provider_rejects_batch_larger_than_workspace() -> None: spec = _spec(tp_size=4) workspace = _cpu_workspace(batch_size=1, head_count=5) diff --git a/tests/test_mla_kernels.py b/tests/test_mla_kernels.py index e59fd278..fb5965c7 100644 --- a/tests/test_mla_kernels.py +++ b/tests/test_mla_kernels.py @@ -342,6 +342,212 @@ def test_mla_decode_matches_torch_matrix( ) +@pytest.mark.parametrize("reduce_heads", [False, True]) +@CUDA_REQUIRED +def test_mla_decode_writes_raw_attention_scores(reduce_heads: bool) -> None: + torch.manual_seed(211) + case = _make_decode_case(batch_size=2, head_count=20, max_context_len=33) + q_latent, q_rope, latent_cache, rope_cache = case[:4] + active_slots, request_indices, context_lens = case[4:] + output = torch.empty_like(q_latent) + score_shape = ( + (2, active_slots.shape[1]) + if reduce_heads + else (2, 20, active_slots.shape[1]) + ) + scores = torch.full( + score_shape, + -1.0e20, + dtype=torch.float32, + device="cuda", + ) + workspace = allocate_mla_decode_workspace( + batch_size=2, + head_count=20, + device="cuda", + ) + + run_mla_decode( + q_latent, + q_rope, + latent_cache, + rope_cache, + active_slots, + request_indices, + context_lens, + output, + workspace, + softmax_scale=GLM_MLA_SOFTMAX_SCALE, + attn_score=scores, + ) + torch.cuda.synchronize() + + for batch_idx in range(2): + length = int(context_lens[batch_idx].item()) + request_row = int(request_indices[batch_idx].item()) + slots = active_slots[request_row, :length].long() + expected = torch.matmul( + q_latent[batch_idx].float(), latent_cache[slots, 0].float().T + ) + torch.matmul( + q_rope[batch_idx].float(), rope_cache[slots, 0].float().T + ) + actual = scores[batch_idx, :length] if reduce_heads else scores[batch_idx, :, :length] + if reduce_heads: + expected = expected.max(dim=0).values + torch.testing.assert_close(actual, expected, rtol=3e-2, atol=3e-2) + assert torch.all(scores[batch_idx, ..., length:] == -1.0e20) + + +@CUDA_REQUIRED +def test_mla_decode_score_capacity_can_be_smaller_than_slot_table() -> None: + torch.manual_seed(219) + case = _make_decode_case(batch_size=1, head_count=20, max_context_len=33) + q_latent, q_rope, latent_cache, rope_cache = case[:4] + active_slots, request_indices, _context_lens = case[4:] + context_lens = torch.tensor([17], dtype=torch.int32, device="cuda") + output = torch.empty_like(q_latent) + scores = torch.empty((1, 17), dtype=torch.float32, device="cuda") + workspace = allocate_mla_decode_workspace( + batch_size=1, + head_count=20, + device="cuda", + ) + + run_mla_decode( + q_latent, + q_rope, + latent_cache, + rope_cache, + active_slots, + request_indices, + context_lens, + output, + workspace, + softmax_scale=GLM_MLA_SOFTMAX_SCALE, + attn_score=scores, + max_context_len=17, + ) + torch.cuda.synchronize() + + assert torch.isfinite(output).all() + assert torch.isfinite(scores).all() + + +@CUDA_REQUIRED +def test_mla_reduced_scores_reset_before_each_decode_step() -> None: + torch.manual_seed(223) + case = _make_decode_case(batch_size=1, head_count=20, max_context_len=33) + q_latent, q_rope, latent_cache, rope_cache = case[:4] + active_slots, request_indices, context_lens = case[4:] + output = torch.empty_like(q_latent) + scores = torch.empty( + (1, active_slots.shape[1]), + dtype=torch.float32, + device="cuda", + ) + workspace = allocate_mla_decode_workspace( + batch_size=1, + head_count=20, + device="cuda", + ) + + run_mla_decode( + q_latent, + q_rope, + latent_cache, + rope_cache, + active_slots, + request_indices, + context_lens, + output, + workspace, + softmax_scale=GLM_MLA_SOFTMAX_SCALE, + attn_score=scores, + ) + first_scores = scores.clone() + run_mla_decode( + torch.zeros_like(q_latent), + torch.zeros_like(q_rope), + latent_cache, + rope_cache, + active_slots, + request_indices, + context_lens, + output, + workspace, + softmax_scale=GLM_MLA_SOFTMAX_SCALE, + attn_score=scores, + ) + torch.cuda.synchronize() + + length = int(context_lens[0].item()) + assert torch.any(first_scores[0, :length] != 0) + torch.testing.assert_close( + scores[0, :length], + torch.zeros_like(scores[0, :length]), + ) + assert torch.all(scores[0, length:] == -1.0e20) + + +@CUDA_REQUIRED +def test_mla_decode_cuda_graph_replay_resets_reduced_scores() -> None: + torch.manual_seed(227) + case = _make_decode_case(batch_size=2, head_count=20, max_context_len=33) + q_latent, q_rope, latent_cache, rope_cache = case[:4] + active_slots, request_indices, context_lens = case[4:] + output = torch.empty_like(q_latent) + scores = torch.empty( + (2, active_slots.shape[1]), + dtype=torch.float32, + device="cuda", + ) + workspace = allocate_mla_decode_workspace( + batch_size=2, + head_count=20, + device="cuda", + ) + + def run_decode() -> None: + run_mla_decode( + q_latent, + q_rope, + latent_cache, + rope_cache, + active_slots, + request_indices, + context_lens, + output, + workspace, + softmax_scale=GLM_MLA_SOFTMAX_SCALE, + attn_score=scores, + validate_metadata=False, + ) + + run_decode() + run_decode() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run_decode() + + q_latent.zero_() + q_rope.zero_() + scores.fill_(12345.0) + graph.replay() + graph_output = output.clone() + graph_scores = scores.clone() + run_decode() + torch.cuda.synchronize() + + torch.testing.assert_close(graph_output, output, rtol=0, atol=0) + torch.testing.assert_close(graph_scores, scores, rtol=0, atol=0) + for batch_idx, length in enumerate(context_lens.tolist()): + torch.testing.assert_close( + graph_scores[batch_idx, :length], + torch.zeros_like(graph_scores[batch_idx, :length]), + ) + assert torch.all(graph_scores[batch_idx, length:] == -1.0e20) + @CUDA_REQUIRED def test_mla_decode_zeroes_padded_rows() -> None: torch.manual_seed(19) diff --git a/tests/test_omnikv_decode_score_lifecycle.py b/tests/test_omnikv_decode_score_lifecycle.py index 4074a4a3..445eb6b0 100644 --- a/tests/test_omnikv_decode_score_lifecycle.py +++ b/tests/test_omnikv_decode_score_lifecycle.py @@ -1,5 +1,5 @@ from types import SimpleNamespace -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import torch @@ -22,6 +22,10 @@ def get_layer_batch_states(self, layer_idx): req_indices=torch.tensor([0], dtype=torch.int32), ) + def get_layer_buffer_req_to_token_slots(self, layer_idx): + del layer_idx + return torch.arange(self.context_len, dtype=torch.int32).reshape(1, -1) + def _make_controller(): layers = 4 @@ -125,3 +129,45 @@ def test_omnikv_graph_reset_and_keepalive_cover_the_shared_workspace(): controller.clear_decode_attn_score_buffers() assert controller._omnikv_decode_attn_score_buffer is None + + +def test_omnikv_decode_graph_reuses_selection_output_buffers(): + controller = _make_controller() + states = controller.layer_batch_sparse_states + + def fake_build(*args, **kwargs): + del args + keep = kwargs["keep_indices_out"] + slots = kwargs["active_slots_out"] + context_lens = kwargs["new_context_lens_out"] + keep.copy_(torch.arange(keep.shape[1], dtype=torch.int32).expand_as(keep)) + slots.copy_(keep) + context_lens.fill_(keep.shape[1]) + return keep, slots, context_lens + + pointers = [] + with patch( + "sparsevllm.engine.sparse_controller.build_omnikv_keep_and_slots", + side_effect=fake_build, + ): + for _ in range(2): + states[0].attn_score = torch.arange(6, dtype=torch.float32).reshape(1, 6) + controller._update_dynamic_omnikv_indices(0, [1]) + pointers.append( + tuple( + int(tensor.data_ptr()) + for tensor in ( + states[1].active_indices, + states[1].active_slots, + states[1].context_lens, + states[1].req_indices, + ) + ) + ) + + assert pointers[0] == pointers[1] + keepalive_ptrs = { + int(tensor.data_ptr()) + for tensor in controller.decode_cuda_graph_keepalive_tensors() + } + assert set(pointers[0]).issubset(keepalive_ptrs) diff --git a/tests/test_prefix_cache.py b/tests/test_prefix_cache.py index ce09cb5a..6d05e4bb 100644 --- a/tests/test_prefix_cache.py +++ b/tests/test_prefix_cache.py @@ -12,7 +12,11 @@ from sparsevllm.config import Config from sparsevllm.engine.cache_manager.quest import QuestCacheManager, QuestPrefixBlockPayload +from sparsevllm.configs.model import RuntimeLayout +from sparsevllm.engine.cache_manager import MlaLatentPayload +from sparsevllm.engine.cache_manager.omnikv import OmniKVCacheManager from sparsevllm.engine.cache_manager.standard import StandardCacheManager, StandardPrefixBlockPayload +from sparsevllm.engine.cache_manager.storage import MlaLatentStorage from sparsevllm.engine.cache_manager.prefix_offload import ( QuestPrefixOffloadController, StandardPrefixOffloadController, @@ -93,11 +97,12 @@ def _make_config(**kwargs): return Config(model=str(model_dir), **kwargs) -def _make_standard_manager_for_prefix(block_size=2): - cfg = _cfg(block_size=block_size) +def _make_standard_manager_for_prefix(block_size=2, method=""): + cfg = _cfg(method=method, block_size=block_size) cfg.num_kvcache_slots = 90 fingerprint = build_prefix_cache_fingerprint(cfg, block_size) - manager = object.__new__(StandardCacheManager) + manager_type = OmniKVCacheManager if method == "omnikv" else StandardCacheManager + manager = object.__new__(manager_type) manager.config = cfg manager.device = torch.device("cpu") manager.enable_prefix_caching = True @@ -1308,6 +1313,7 @@ def test_resolve_prefix_cache_block_size_uses_quest_page_size(): def test_prefix_cache_supported_method_allowlist(): assert PREFIX_CACHE_SUPPORTED_METHODS == { "", + "streamingllm", "omnikv", "quest", "snapkv", @@ -1351,12 +1357,13 @@ def test_config_resolves_prefix_cache_defaults(): assert cfg.resolved_prefix_cache_mode == "disabled" -def test_config_rejects_unsupported_prefix_cache_methods(): - with pytest.raises(ValueError, match="not supported"): - _make_config( - vllm_sparse_method="streamingllm", - enable_prefix_caching=True, - ) +def test_config_accepts_streamingllm_chain_prefix_cache(): + cfg = _make_config( + vllm_sparse_method="streamingllm", + enable_prefix_caching=True, + ) + + assert cfg.resolved_prefix_cache_mode == "chain" def test_config_rejects_unvalidated_prefix_cache_options(): @@ -1507,6 +1514,170 @@ def test_standard_attach_pins_prefix_slots_and_free_seq_keeps_cached_slots(): assert manager.seq_id_to_row == {} +def test_standard_latent_prefix_restores_mla_payload_and_cleans_request_state(): + manager = _make_standard_manager_for_prefix(block_size=2) + storage = MlaLatentStorage( + kv_lora_rank=512, + rope_dim=64, + dtype=torch.bfloat16, + ) + storage.allocate(num_layers=1, num_slots=100, device=torch.device("cpu")) + manager.attention_cache_storage = storage + + owner = Sequence([1, 2]) + owner_slots = manager._allocate(owner.seq_id, 2).clone() + assert storage.latent_cache is not None + assert storage.rope_cache is not None + storage.latent_cache[0, owner_slots[0]].fill_(11) + storage.latent_cache[0, owner_slots[1]].fill_(22) + storage.rope_cache[0, owner_slots[0]].fill_(33) + storage.rope_cache[0, owner_slots[1]].fill_(44) + manager._record_prefix_materialization(owner, [1, 2], owner_slots) + manager.on_forward_end([owner], is_prefill=True) + manager.free_seq(owner.seq_id) + + replay = Sequence([1, 2, 9]) + manager.refresh_prefix_cache_hit(replay) + assert replay.prefix_cache_hit_len == 2 + manager._attach_prefix_cache_if_needed(replay) + replay_row = manager.seq_id_to_row[replay.seq_id] + replay_slots = manager.buffer_req_to_token_slots[replay_row, :2].clone() + assert replay_slots.tolist() == owner_slots.tolist() + payload = storage.layer_payload(0) + assert isinstance(payload, MlaLatentPayload) + torch.testing.assert_close( + payload.latent_cache[replay_slots, 0, 0], + torch.tensor([11, 22], dtype=torch.bfloat16), + ) + torch.testing.assert_close( + payload.rope_cache[replay_slots, 0, 0], + torch.tensor([33, 44], dtype=torch.bfloat16), + ) + + manager.free_seq(replay.seq_id) + assert replay.seq_id not in manager.seq_id_to_row + assert replay.seq_id not in manager.seq_id_to_prefix_blocks + assert replay.seq_id not in manager.seq_id_to_cached_ranges + assert replay.seq_id not in manager.prefix_runtime_states + manager.reset_prefix_cache() + assert len(manager.prefix_cache) == 0 + assert manager._num_free_slots == 90 + + unrelated = Sequence([7, 8, 9]) + manager.refresh_prefix_cache_hit(unrelated) + assert unrelated.prefix_cache_hit_len == 0 + assert unrelated.seq_id not in manager.seq_id_to_row + + +def _assert_standard_latent_prefix_full_lifecycle(method: str): + manager = _make_standard_manager_for_prefix(block_size=2, method=method) + manager.max_model_len = 16 + manager.num_layers = 1 + manager.num_kv_layers = 1 + manager.runtime_layout = RuntimeLayout.dense(1) + storage = MlaLatentStorage( + kv_lora_rank=512, + rope_dim=64, + dtype=torch.bfloat16, + ) + storage.allocate(num_layers=1, num_slots=100, device=torch.device("cpu")) + manager.attention_cache_storage = storage + + owner = Sequence([1, 2, 9]) + owner.current_chunk_size = 3 + input_ids, positions, cu_seqlens_q = manager.prepare_step([owner], is_prefill=True) + assert input_ids.tolist() == [1, 2, 9] + assert positions.tolist() == [0, 1, 2] + assert cu_seqlens_q.tolist() == [0, 3] + owner_slots = manager.layer_batch_state.slot_mapping.clone() + payload = storage.layer_payload(0) + payload.latent_cache[owner_slots] = ( + torch.tensor([11, 22, 99], dtype=torch.bfloat16) + .view(3, 1, 1) + .expand(3, 1, 512) + ) + payload.rope_cache[owner_slots] = ( + torch.tensor([33, 44, 88], dtype=torch.bfloat16) + .view(3, 1, 1) + .expand(3, 1, 64) + ) + manager.on_forward_end([owner], is_prefill=True) + manager.free_seq(owner.seq_id) + assert len(manager.prefix_cache) == 1 + assert manager._num_free_slots == 88 + + replay = Sequence([1, 2, 8]) + manager.refresh_prefix_cache_hit(replay) + assert replay.prefix_cache_hit_len == 2 + replay.num_prefilled_tokens = replay.prefix_cache_hit_len + replay.current_chunk_size = 1 + input_ids, positions, cu_seqlens_q = manager.prepare_step([replay], is_prefill=True) + assert input_ids.tolist() == [8] + assert positions.tolist() == [2] + assert cu_seqlens_q.tolist() == [0, 1] + replay_row = manager.seq_id_to_row[replay.seq_id] + replay_slots = manager.buffer_req_to_token_slots[replay_row, :3].clone() + assert replay_slots[:2].tolist() == owner_slots[:2].tolist() + torch.testing.assert_close( + payload.latent_cache[replay_slots[:2], 0, 0], + torch.tensor([11, 22], dtype=torch.bfloat16), + ) + torch.testing.assert_close( + payload.rope_cache[replay_slots[:2], 0, 0], + torch.tensor([33, 44], dtype=torch.bfloat16), + ) + payload.latent_cache[replay_slots[2:]] = 77 + payload.rope_cache[replay_slots[2:]] = 66 + manager.on_forward_end([replay], is_prefill=True) + manager.free_seq(replay.seq_id) + assert replay.seq_id not in manager.seq_id_to_row + assert replay.seq_id not in manager.seq_id_to_prefix_blocks + assert replay.seq_id not in manager.seq_id_to_cached_ranges + assert replay.seq_id not in manager.seq_id_to_materialized_blocks + assert replay.seq_id not in manager.pending_prefix_blocks + assert replay.seq_id not in manager.prefix_runtime_states + assert manager._num_free_slots == 88 + + deleted = manager.prefix_cache_delete_subtree([1, 2]) + assert deleted["deleted_block_count"] == 1 + assert len(manager.prefix_cache) == 0 + assert manager._num_free_slots == 90 + + replacement = Sequence([7, 8]) + replacement.current_chunk_size = 2 + manager.prepare_step([replacement], is_prefill=True) + replacement_slots = manager.layer_batch_state.slot_mapping.clone() + assert replacement_slots.tolist() == owner_slots[:2].tolist() + payload.latent_cache[replacement_slots] = 55 + payload.rope_cache[replacement_slots] = 44 + torch.testing.assert_close( + payload.latent_cache[replacement_slots, 0, 0], + torch.tensor([55, 55], dtype=torch.bfloat16), + ) + torch.testing.assert_close( + payload.rope_cache[replacement_slots, 0, 0], + torch.tensor([44, 44], dtype=torch.bfloat16), + ) + manager.on_forward_end([replacement], is_prefill=True) + manager.free_seq(replacement.seq_id) + manager.prefix_cache_delete_subtree([7, 8]) + assert manager._num_free_slots == 90 + assert manager.seq_id_to_row == {} + assert manager.seq_id_to_prefix_blocks == {} + assert manager.seq_id_to_cached_ranges == {} + assert manager.seq_id_to_materialized_blocks == {} + assert manager.pending_prefix_blocks == {} + assert manager.prefix_runtime_states == {} + + +def test_standard_latent_prefix_full_lifecycle_restores_and_reuses_slots(): + _assert_standard_latent_prefix_full_lifecycle("") + + +def test_omnikv_latent_prefix_full_lifecycle_restores_and_reuses_slots(): + _assert_standard_latent_prefix_full_lifecycle("omnikv") + + def test_standard_offload_gpu_pressure_only_demotes_dual_resident_blocks(): manager = _make_standard_manager_for_prefix(block_size=2) controller = _FakePrefixOffloadController(manager.prefix_cache) diff --git a/tests/test_prefix_cache_bench.py b/tests/test_prefix_cache_bench.py index 3fa049ff..f813c099 100644 --- a/tests/test_prefix_cache_bench.py +++ b/tests/test_prefix_cache_bench.py @@ -185,6 +185,15 @@ def test_prefix_cache_bench_exposes_h2o_chain_case(): assert bench.CASE_ALIASES["h2o"] == "chain_h2o" +def test_prefix_cache_bench_exposes_streamingllm_chain_case(): + assert bench.CASE_PRESETS["chain_streamingllm"] == { + "method": "streamingllm", + "enable_prefix_caching": True, + "label": "StreamingLLM, linear chain prefix cache", + } + assert bench.CASE_ALIASES["streamingllm"] == "chain_streamingllm" + + def test_prefix_cache_bench_labels_chain_reuse_as_logical_tokens(): summary = bench._summarize_records( case_name="chain_snapkv", @@ -247,6 +256,8 @@ def test_prefix_cache_bench_engine_kwargs_are_sparsevllm_config_fields(): args = types.SimpleNamespace( gpu_memory_utilization=0.65, tensor_parallel_size=1, + expert_parallel_size=2, + decode_cuda_graph=True, max_active_requests=4, max_num_batched_tokens=8192, chunk_prefill_size=4096, @@ -270,6 +281,89 @@ def test_prefix_cache_bench_engine_kwargs_are_sparsevllm_config_fields(): config_fields = set(Config.__dataclass_fields__) unknown = sorted(set(normalized.infer_config) - config_fields) assert unknown == [] + assert normalized.infer_config["expert_parallel_size"] == 2 + assert normalized.infer_config["decode_cuda_graph"] is True + + +def test_prefix_cache_bench_requires_graph_capture_and_replay_on_every_rank(): + args = _summary_args() + args.decode_cuda_graph = True + summary = bench._summarize_records( + case_name="prefix_full", + case_config=bench.CASE_PRESETS["prefix_full"], + records=[], + args=args, + engine_kwargs={}, + cache_stats_before={}, + cache_stats_after={}, + peak_memory_gb=0.0, + elapsed_s=1.0, + decode_graph_before=[ + { + "world_rank": rank, + "capture_count": 0, + "replay_count": 0, + "eager_static_count": 0, + "force_eager_count": 0, + } + for rank in (0, 1) + ], + decode_graph_after=[ + { + "world_rank": rank, + "capture_count": 1, + "replay_count": 3, + "eager_static_count": 0, + "force_eager_count": 0, + } + for rank in (0, 1) + ], + ) + + assert summary["decode_cuda_graph_failures"] == [] + assert summary["decode_cuda_graph_delta"] == [ + { + "world_rank": rank, + "capture_count": 1, + "replay_count": 3, + "eager_static_count": 0, + "force_eager_count": 0, + } + for rank in (0, 1) + ] + + +def test_prefix_cache_bench_accepts_warmup_capture_and_business_replay(): + args = _summary_args() + args.decode_cuda_graph = True + counters = { + "capture_count": 1, + "eager_static_count": 0, + "force_eager_count": 0, + } + + summary = bench._summarize_records( + case_name="prefix_full", + case_config=bench.CASE_PRESETS["prefix_full"], + records=[], + args=args, + engine_kwargs={}, + cache_stats_before={}, + cache_stats_after={}, + peak_memory_gb=0.0, + elapsed_s=1.0, + decode_graph_before=[ + {"world_rank": 0, "replay_count": 1, **counters} + ], + decode_graph_after=[ + {"world_rank": 0, "replay_count": 5, **counters} + ], + ) + + assert summary["status"] == "success" + assert summary["decode_cuda_graph_failures"] == [] + assert summary["decode_cuda_graph_delta"][0]["capture_count"] == 0 + assert summary["decode_cuda_graph_delta"][0]["replay_count"] == 4 def test_prefix_cache_bench_trace_uses_resolved_sparse_budgets(): diff --git a/tests/test_rkv_skipkv_methods.py b/tests/test_rkv_skipkv_methods.py index 1f6d6789..7070d825 100644 --- a/tests/test_rkv_skipkv_methods.py +++ b/tests/test_rkv_skipkv_methods.py @@ -4,6 +4,7 @@ import unittest from unittest.mock import patch +import numpy as np import torch from sparsevllm.config import Config, RuntimeLayout @@ -14,6 +15,7 @@ PrefillComputeView, ) from sparsevllm.engine.cache_manager.rkv import RKVCacheManager +from sparsevllm.engine.cache_manager.storage import MlaLatentStorage from sparsevllm.engine.cache_manager.skipkv import ( SkipKVCacheManager, SkipKVSentence, @@ -21,6 +23,10 @@ ) from sparsevllm.engine.activation_controller import ActivationController from sparsevllm.engine.sequence import Sequence +from sparsevllm.engine.sparse_controller import ( + LayerBatchSparseState, + SparseController, +) from sparsevllm.method_registry import ( get_default_prefill_schedule_policy, normalize_sparse_method, @@ -99,6 +105,19 @@ def test_rkv_joint_retention_score_uses_paper_lambda(self): expected = 0.25 * importance - 0.75 * redundancy self.assertTrue(torch.allclose(score, expected)) + def test_attention_key_materializer_registration_is_idempotent_and_strict(self): + manager = object.__new__(RKVCacheManager) + manager.runtime_layout = RuntimeLayout.dense(1) + first = lambda view: view.payload + second = lambda view: view.payload + + manager.register_attention_key_materializer(0, first) + manager.register_attention_key_materializer(0, first) + + self.assertTrue(manager.has_attention_key_materializer(0)) + with self.assertRaisesRegex(RuntimeError, "already bound"): + manager.register_attention_key_materializer(0, second) + def test_skipkv_segment_penalty_marks_older_similar_segment(self): keys = torch.tensor( [ @@ -544,6 +563,313 @@ def test_rkv_query_cache_tracks_observation_tokens_not_interval(self): expected = torch.cat((q_prefill[3:5], q_decode), dim=0) torch.testing.assert_close(manager._rkv_query_cache[0][0, cols], expected) + def test_rkv_mla_query_attention_scores_match_expanded_key_oracle(self): + torch.manual_seed(37) + kv_len = 6 + observation_tokens = 2 + num_heads = 2 + seq = Sequence([1]) + manager = object.__new__(RKVCacheManager) + manager.runtime_layout = RuntimeLayout.dense(1) + manager.config = SimpleNamespace( + rkv_observation_tokens=observation_tokens, + sparse_attn_score_dtype="float32", + ) + manager.device = torch.device("cpu") + manager._rkv_query_cache_enabled = True + manager._rkv_observation_tokens = observation_tokens + manager.seq_id_to_row = [{seq.seq_id: 0}] + slots = torch.tensor([5, 1, 8, 3, 7, 2], dtype=torch.int32) + manager.buffer_req_to_token_slots = [slots.view(1, kv_len)] + storage = MlaLatentStorage( + kv_lora_rank=512, + rope_dim=64, + dtype=torch.bfloat16, + ) + storage.allocate(num_layers=1, num_slots=10, device=torch.device("cpu")) + manager.attention_cache_storage = storage + payload = storage.layer_payload(0) + payload.latent_cache.copy_( + torch.randn_like(payload.latent_cache) + ) + payload.rope_cache.copy_(torch.randn_like(payload.rope_cache)) + k_weight = torch.randn( + num_heads, + 192, + 512, + dtype=torch.bfloat16, + ) + + def materialize(view): + flat_slots = view.active_slots.long().reshape(-1) + latent = view.payload.latent_cache[flat_slots, 0] + rope = view.payload.rope_cache[flat_slots, 0] + k_nope = torch.einsum( + "cr,hdr->chd", + latent.float(), + k_weight.float(), + ).to(torch.bfloat16) + keys = torch.cat( + (k_nope, rope[:, None, :].expand(-1, num_heads, -1)), + dim=-1, + ) + return keys.view(*view.active_slots.shape, num_heads, 256) + + manager.register_attention_key_materializer(0, materialize) + manager._rkv_query_cache = [ + torch.empty( + (1, observation_tokens, num_heads, 256), + dtype=torch.bfloat16, + ) + ] + manager._rkv_query_positions = [ + torch.full((1, observation_tokens), -1, dtype=torch.int32) + ] + query_positions = torch.tensor([4, 5], dtype=torch.long) + query_cols = query_positions.remainder(observation_tokens) + q_window = torch.randn( + observation_tokens, + num_heads, + 256, + dtype=torch.bfloat16, + ) + manager._rkv_query_cache[0][0, query_cols] = q_window + manager._rkv_query_positions[0][0, query_cols] = query_positions.to( + torch.int32 + ) + + scores = manager.rkv_query_attention_scores( + 0, + seq, + kv_len, + candidate_start=1, + num_recent_tokens=1, + ) + + candidate_positions = torch.arange(1, 5, dtype=torch.long) + candidate_slots = slots.index_select(0, candidate_positions).long() + latent = payload.latent_cache[candidate_slots, 0] + rope = payload.rope_cache[candidate_slots, 0] + k_nope = torch.einsum( + "cr,hdr->chd", + latent.float(), + k_weight.float(), + ).to(torch.bfloat16).float() + expanded_keys = torch.cat( + (k_nope, rope.float()[:, None, :].expand(-1, num_heads, -1)), + dim=-1, + ) + logits = torch.einsum( + "qhd,chd->qhc", + q_window.float(), + expanded_keys, + ) + logits.mul_(256**-0.5) + valid = candidate_positions.unsqueeze(0) <= query_positions.unsqueeze(1) + logits.masked_fill_(~valid[:, None, :], torch.finfo(logits.dtype).min) + expected = torch.zeros((kv_len,), dtype=torch.float32) + expected[1:5] = ( + torch.softmax(logits, dim=-1) + .masked_fill(~valid[:, None, :], 0.0) + .mean(dim=0) + .amax(dim=0) + ) + torch.testing.assert_close(scores, expected, rtol=1e-4, atol=1e-4) + + def test_rkv_mla_redundancy_uses_actual_expanded_keys(self): + manager = object.__new__(RKVCacheManager) + manager.runtime_layout = RuntimeLayout.dense(1) + manager.config = SimpleNamespace( + num_sink_tokens=1, + num_recent_tokens=1, + rkv_redundancy_window=0, + rkv_similarity_threshold=0.0, + rkv_recent_similar_keep=0, + rkv_max_redundancy_tokens=8, + rkv_alpha=0.5, + ) + seq = Sequence([1]) + manager.seq_id_to_row = [{seq.seq_id: 0}] + manager.buffer_req_to_token_slots = [ + torch.arange(6, dtype=torch.int32).view(1, 6) + ] + storage = MlaLatentStorage( + kv_lora_rank=512, + rope_dim=64, + dtype=torch.bfloat16, + ) + storage.allocate(num_layers=1, num_slots=6, device=torch.device("cpu")) + manager.attention_cache_storage = storage + payload = storage.layer_payload(0) + for slot in range(6): + payload.latent_cache[slot].fill_(slot) + payload.rope_cache[slot].fill_(slot + 100) + num_heads = 2 + k_weight = torch.arange( + num_heads * 192 * 512, + dtype=torch.float32, + ).remainder(17).sub_(8).mul_(1.0e-4).view(num_heads, 192, 512) + + def materialize(view): + flat_slots = view.active_slots.long().reshape(-1) + latent = view.payload.latent_cache[flat_slots, 0] + rope = view.payload.rope_cache[flat_slots, 0] + k_nope = torch.einsum( + "cr,hdr->chd", + latent.float(), + k_weight, + ).to(torch.bfloat16) + keys = torch.cat( + (k_nope, rope[:, None, :].expand(-1, num_heads, -1)), + dim=-1, + ) + return keys.view(*view.active_slots.shape, num_heads, 256) + + manager.register_attention_key_materializer(0, materialize) + captured = [] + + def fake_redundancy(keys, **kwargs): + del kwargs + captured.append(keys.clone()) + return torch.zeros((keys.shape[0],), dtype=torch.float32) + + with patch.object( + manager, + "redundancy_scores_from_keys", + side_effect=fake_redundancy, + ): + manager.select_rkv_indices( + 0, + seq, + torch.arange(6, dtype=torch.float32), + kv_len=6, + budget=4, + ) + + self.assertEqual(len(captured), 1) + self.assertEqual(tuple(captured[0].shape), (4, num_heads, 256)) + expected_slots = torch.arange(1, 5, dtype=torch.int32) + expected_view = manager.build_attention_key_compute_view( + 0, + expected_slots, + ) + torch.testing.assert_close(captured[0], materialize(expected_view)) + + def test_rkv_mla_budget_trigger_compacts_actual_latent_slots(self): + row_len = 6 + storage = MlaLatentStorage( + kv_lora_rank=512, + rope_dim=64, + dtype=torch.bfloat16, + ) + storage.allocate(num_layers=1, num_slots=8, device=torch.device("cpu")) + manager = object.__new__(RKVCacheManager) + manager.attention_cache_storage = storage + manager.kv_cache = None + manager.device = torch.device("cpu") + manager.runtime_layout = RuntimeLayout.dense(1) + manager.num_layers = 1 + manager.num_kv_layers = 1 + manager._uniform_decode_metadata = True + manager.buffer_req_to_token_slots_tensor = torch.zeros( + (1, 1, 8), + dtype=torch.int32, + ) + manager.buffer_req_to_token_slots_tensor[0, 0, :row_len] = torch.arange( + row_len, + dtype=torch.int32, + ) + manager.buffer_req_to_token_slots = [ + manager.buffer_req_to_token_slots_tensor[0] + ] + manager.seq_id_to_row = [{0: 0}] + manager.row_seq_lens = [np.asarray([row_len], dtype=np.int32)] + manager.free_slots_stack_tensor = None + manager.free_slots_stack = [torch.zeros((8,), dtype=torch.int32)] + manager._num_free_slots = [0] + manager._rkv_query_cache_enabled = True + manager._rkv_observation_tokens = 2 + manager._rkv_batch_clear_query_cache_rows = True + manager._rkv_query_cache = [ + torch.zeros((1, 2, 2, 256), dtype=torch.bfloat16) + ] + manager._rkv_query_positions = [ + torch.full((1, 2), -1, dtype=torch.int32) + ] + query_positions = torch.tensor([4, 5], dtype=torch.long) + query_cols = query_positions.remainder(2) + manager._rkv_query_cache[0][0, query_cols, :, 0] = 1 + manager._rkv_query_positions[0][0, query_cols] = query_positions.to( + torch.int32 + ) + manager.config = SimpleNamespace( + num_sink_tokens=1, + decode_keep_tokens=2, + num_recent_tokens=1, + rkv_compression_interval=2, + rkv_redundancy_window=4, + rkv_similarity_threshold=0.0, + rkv_recent_similar_keep=0, + rkv_max_redundancy_tokens=8, + rkv_alpha=1.0, + sparse_attn_score_dtype="float32", + ) + + payload = storage.layer_payload(0) + key_values = torch.tensor([0, -2, 5, 1, 4, 0], dtype=torch.bfloat16) + for slot, key_value in enumerate(key_values): + payload.latent_cache[slot].zero_() + payload.latent_cache[slot, 0, 0] = key_value + payload.rope_cache[slot].fill_(slot + 100) + + def materialize(view): + flat_slots = view.active_slots.long().reshape(-1) + latent = view.payload.latent_cache[flat_slots, 0] + rope = view.payload.rope_cache[flat_slots, 0] + keys = torch.zeros( + (flat_slots.numel(), 2, 256), + dtype=torch.bfloat16, + ) + keys[:, :, 0] = latent[:, None, 0] + keys[:, :, 192:] = rope[:, None, :] + return keys.view(*view.active_slots.shape, 2, 256) + + manager.register_attention_key_materializer(0, materialize) + seq = Sequence(list(range(row_len))) + seq.seq_id = 0 + controller = object.__new__(SparseController) + controller.cache_manager = manager + controller.device = torch.device("cpu") + controller.sparse_method = "rkv" + controller.num_layers = 1 + controller.num_sink = 1 + controller.num_recent = 1 + controller.decode_keep_tokens = 2 + controller.config = manager.config + controller.layer_batch_sparse_states = { + 0: LayerBatchSparseState( + context_lens=torch.tensor([row_len], dtype=torch.int32) + ) + } + controller._is_kv_layer = lambda layer_idx: int(layer_idx) == 0 + + controller._rkv_decode_eviction([seq]) + + active_slots = manager.buffer_req_to_token_slots[0][0, :4].long() + self.assertEqual(active_slots.tolist(), [0, 2, 4, 5]) + self.assertEqual(manager.row_seq_lens[0].tolist(), [4]) + self.assertEqual(sorted(manager.free_slots_stack[0][:2].tolist()), [1, 3]) + self.assertEqual(manager._num_free_slots, [2]) + torch.testing.assert_close( + payload.latent_cache[active_slots, 0, 0], + torch.tensor([0, 5, 4, 0], dtype=torch.bfloat16), + ) + torch.testing.assert_close( + payload.rope_cache[active_slots, 0, 0], + torch.tensor([100, 102, 104, 105], dtype=torch.bfloat16), + ) + self.assertEqual(manager._rkv_query_positions[0][0].tolist(), [-1, -1]) + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for R-KV query score kernel tests.") def test_rkv_query_attention_scores_reuse_prefill_score_kernel(self): from tests.test_prefill_score_kernel import _prefill_score_baseline @@ -611,6 +937,32 @@ def test_rkv_query_attention_scores_reuse_prefill_score_kernel(self): )[0, :kv_len] torch.testing.assert_close(scores, expected, rtol=2e-2, atol=2e-2) + candidate_positions = torch.arange( + candidate_start, + kv_len - num_recent_tokens, + dtype=torch.long, + device=device, + ) + candidate_slots = manager.buffer_req_to_token_slots[0][ + 0, + candidate_start : kv_len - num_recent_tokens, + ].long() + materialized = torch.zeros_like(scores) + materialized[candidate_start : kv_len - num_recent_tokens] = ( + RKVCacheManager.attention_scores_from_materialized_keys( + q_window, + k_cache.index_select(0, candidate_slots), + positions, + candidate_positions, + ) + ) + torch.testing.assert_close( + materialized, + scores, + rtol=2e-2, + atol=2e-2, + ) + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for R-KV query score kernel tests.") def test_rkv_query_attention_scores_batch_matches_single_path(self): torch.manual_seed(31) diff --git a/tests/test_sparse_state_summary.py b/tests/test_sparse_state_summary.py index 5e671e92..151c8b68 100644 --- a/tests/test_sparse_state_summary.py +++ b/tests/test_sparse_state_summary.py @@ -1,9 +1,11 @@ +import os from types import SimpleNamespace from unittest.mock import patch import torch from sparsevllm.engine.cache_manager.base import _debug_tensor_summary +from sparsevllm.distributed import ParallelContext, ParallelGroup from sparsevllm.engine.model_runner import ModelRunner from sparsevllm.engine.sparse_controller import LayerBatchSparseState, SparseController @@ -61,14 +63,44 @@ def test_model_runner_gathers_one_debug_summary_per_world_rank(): runner = object.__new__(ModelRunner) runner.world_size = 2 runner.rank = 0 - runner.parallel_context = SimpleNamespace( - world_rank=0, - ep_rank=0, - world=SimpleNamespace(process_group=object()), + process_group = object() + world_group = ParallelGroup( + process_group=process_group, + ranks=(0, 1), + rank=0, + size=2, + ) + singleton_group = ParallelGroup( + process_group=None, + ranks=(0,), + rank=0, + size=1, + ) + runner.parallel_context = ParallelContext( + world=world_group, + tensor=singleton_group, + expert=world_group, + data=singleton_group, + moe_tensor=singleton_group, ) runner.sparse_controller = SimpleNamespace( debug_state_summary=lambda: {"sparse_method": "", "layers": {}} ) + runner.config = SimpleNamespace(decode_cuda_graph=True) + runner.decode_cuda_graph_runner = SimpleNamespace( + capture_count=1, + replay_count=3, + eager_static_count=0, + force_eager_count=0, + _graphs={"graph": object()}, + last_state_key=SimpleNamespace( + method="snapkv", + batch_size=2, + context_capacity=1024, + is_long_text=True, + capture_sampling=False, + ), + ) def gather(output, local, group): assert group is runner.parallel_context.world.process_group @@ -82,4 +114,78 @@ def gather(output, local, group): assert [summary["world_rank"] for summary in summaries] == [0, 1] assert summaries[0]["state"] == summaries[1]["state"] + assert summaries[0]["decode_cuda_graph"] == { + "enabled": True, + "capture_count": 1, + "replay_count": 3, + "eager_static_count": 0, + "force_eager_count": 0, + "cached_graph_count": 1, + "last_state_key": { + "method": "snapkv", + "batch_size": 2, + "context_capacity": 1024, + "is_long_text": True, + "capture_sampling": False, + }, + } sync_status.assert_called_once_with("debug_sparse_state_summaries", None) + + +def test_tp_debug_replica_consistency_marks_vocab_sharded_logits_not_applicable(): + runner = object.__new__(ModelRunner) + runner.world_size = 2 + runner.parallel_context = ParallelContext( + world=ParallelGroup(process_group=object(), ranks=(0, 1), rank=1, size=2), + tensor=ParallelGroup(process_group=object(), ranks=(0, 1), rank=1, size=2), + expert=ParallelGroup(process_group=None, ranks=(1,), rank=0, size=1), + data=ParallelGroup(process_group=None, ranks=(1,), rank=0, size=1), + ) + runner.model = SimpleNamespace(model=SimpleNamespace(layers=())) + + consistency = runner.debug_replica_consistency() + + assert consistency == { + "last_logits_max_abs": None, + "last_logits_tolerance_ratio": None, + "last_logits_comparison": "not_applicable_tp_vocab_sharded", + "moe_layers": {}, + } + + +def test_nonzero_rank_debug_logits_rpc_returns_none_before_tensor_access(): + runner = object.__new__(ModelRunner) + runner.rank = 1 + + assert runner.debug_last_logits_cpu() is None + + +def test_run_model_does_not_capture_non_tensor_tp_logits(): + runner = object.__new__(ModelRunner) + + class FakeModel: + def __call__(self, input_ids, positions): + return torch.ones(1) + + def compute_logits(self, hidden_states): + return None + + runner.model = FakeModel() + with patch.dict(os.environ, {"SPARSEVLLM_DEBUG_RUNTIME": "1"}): + logits = runner.run_model(torch.ones(1), torch.ones(1), is_prefill=False) + + assert logits is None + assert not hasattr(runner, "debug_last_logits") + + +def test_debug_logits_can_be_refreshed_after_cuda_graph_replay(): + runner = object.__new__(ModelRunner) + capture_value = torch.tensor([[1.0, 2.0]]) + replay_value = torch.tensor([[3.0, 4.0]]) + + with patch.dict(os.environ, {"SPARSEVLLM_DEBUG_RUNTIME": "1"}): + runner._record_debug_logits(capture_value) + runner._record_debug_logits(replay_value) + + torch.testing.assert_close(runner.debug_last_logits, replay_value) + assert runner.debug_last_logits.data_ptr() != replay_value.data_ptr() diff --git a/tests/test_sparsevllm_regression_grading.py b/tests/test_sparsevllm_regression_grading.py index a2fe3317..c0d89c26 100644 --- a/tests/test_sparsevllm_regression_grading.py +++ b/tests/test_sparsevllm_regression_grading.py @@ -48,7 +48,9 @@ H2O_SUPPORTED_MODEL_TYPES, PREFILL_POLICY_LONG_BS1FULL_SHORT_BATCH, get_default_prefill_schedule_policy, + validate_model_runtime_compatibility, ) +from sparsevllm.distributed import ParallelMode, ParallelTopology def _single_process_parallel_context() -> ParallelContext: @@ -196,9 +198,28 @@ def test_h2o_manifest_declares_supported_models_tp_runtime_matrix(self): ["qwen2", "qwen3", "qwen3_moe", "qwen3_5", "qwen3_5_moe", "llama", "minimax_m2"], ) self.assertEqual( - set(method["supported_model_families"]), - set(H2O_SUPPORTED_MODEL_TYPES), + set(H2O_SUPPORTED_MODEL_TYPES) - set(method["supported_model_families"]), + {"glm4_moe_lite"}, ) + for tp_size, ep_size in ( + (1, 1), + (1, 2), + (1, 4), + (2, 1), + (2, 2), + (4, 1), + (4, 2), + (4, 4), + ): + mode = ParallelMode.OUTER_TP_MOE if tp_size > 1 else ParallelMode.STANDARD + validate_model_runtime_compatibility( + model_type="glm4_moe_lite", + sparse_method="h2o", + topology=ParallelTopology(tp_size, ep_size, 1, mode), + enforce_eager=True, + decode_cuda_graph=True, + enable_prefix_caching=True, + ) self.assertEqual(method["supported_tensor_parallel_sizes"], [1, 2]) self.assertEqual(method["performance"]["minimum_prefill_speedup"], 1.0) self.assertIsNone( diff --git a/tests/test_sparsevllm_regression_suite.py b/tests/test_sparsevllm_regression_suite.py index 9af8997e..52317bd6 100644 --- a/tests/test_sparsevllm_regression_suite.py +++ b/tests/test_sparsevllm_regression_suite.py @@ -79,24 +79,34 @@ def test_quality_command_can_enable_tp_prefix_graph_for_supported_methods(self): self.assertEqual(hyper_params["prefix_cache_block_size"], 16) self.assertFalse(hyper_params["decode_cuda_graph_capture_sampling"]) - def test_quality_command_rejects_prefix_cache_for_unsupported_methods(self): - with self.assertRaisesRegex(ValueError, "enable_prefix_caching"): - run_suite._quality_command( - model_id="qwen25_7b", - method_id="streamingllm", - model={"model_path": "/models/qwen", "tokenizer_path": "/models/qwen"}, - method={ - "sparse_method": "streamingllm", - "config": {"sparse_method": "streamingllm"}, - }, - quality={**self._quality_cfg(), "enable_prefix_caching": True}, - performance={ - "decode_cuda_graph": True, - "enforce_eager": False, - "tensor_parallel_size": 2, - }, - output_root=Path("/tmp/sparsevllm-quality"), - ) + def test_quality_command_enables_streamingllm_prefix_graph(self): + cmd = run_suite._quality_command( + model_id="qwen25_7b", + method_id="streamingllm", + model={"model_path": "/models/qwen", "tokenizer_path": "/models/qwen"}, + method={ + "sparse_method": "streamingllm", + "config": {"sparse_method": "streamingllm"}, + }, + quality={ + **self._quality_cfg(), + "enable_prefix_caching": True, + "prefix_cache_block_size": 16, + }, + performance={ + "decode_cuda_graph": True, + "enforce_eager": False, + "tensor_parallel_size": 2, + }, + output_root=Path("/tmp/sparsevllm-quality"), + ) + + hyper_params = json.loads(cmd[cmd.index("--hyper_param") + 1]) + self.assertEqual(hyper_params["tensor_parallel_size"], 2) + self.assertTrue(hyper_params["decode_cuda_graph"]) + self.assertTrue(hyper_params["enable_prefix_caching"]) + self.assertEqual(hyper_params["prefix_cache_block_size"], 16) + self.assertFalse(hyper_params["decode_cuda_graph_capture_sampling"]) def test_perf_command_uses_tp_decode_graph_hyper_params(self): cmd = run_suite._perf_command( diff --git a/tests/test_validate_glm_checkpoint_parallel.py b/tests/test_validate_glm_checkpoint_parallel.py new file mode 100644 index 00000000..7e6ee4b8 --- /dev/null +++ b/tests/test_validate_glm_checkpoint_parallel.py @@ -0,0 +1,198 @@ +from __future__ import annotations + +import pytest + +from scripts.validation.validate_glm_checkpoint_parallel import ( + _validate_parallel_evidence, +) + + +def _graph_summary(enabled: bool = False) -> dict: + return { + "enabled": enabled, + "capture_count": 1 if enabled else 0, + "replay_count": 2 if enabled else 0, + "eager_static_count": 0, + "force_eager_count": 0, + } + + +def _ep4_summary(world_rank: int, local_hit_count: int) -> dict: + layer_key = "46" + local_width = 16 + return { + "world_rank": world_rank, + "parallel": { + "configured": { + "tensor_parallel_size": 1, + "expert_parallel_size": 4, + "data_parallel_size": 1, + "world_size": 4, + }, + "effective": { + "world": { + "rank": world_rank, + "size": 4, + "ranks": [0, 1, 2, 3], + }, + "attention": { + "rank": 0, + "size": 1, + "ranks": [world_rank], + }, + "expert": { + "rank": world_rank, + "size": 4, + "ranks": [0, 1, 2, 3], + }, + "moe_tensor": { + "rank": 0, + "size": 1, + "ranks": [world_rank], + }, + "data": { + "rank": 0, + "size": 1, + "ranks": [world_rank], + }, + }, + "attention_replicated_for_ep": True, + }, + "moe_local": { + layer_key: { + "local_expert_start": world_rank * local_width, + "local_expert_end": (world_rank + 1) * local_width, + "local_hit_count": local_hit_count, + } + }, + "replica_consistency": { + "last_logits_comparison": "compared", + "last_logits_tolerance_ratio": 0.0, + "moe_layers": { + layer_key: { + "topk_ids_mismatch": False, + "topk_weights_tolerance_ratio": 0.0, + "output_tolerance_ratio": 0.0, + } + }, + }, + "decode_cuda_graph": _graph_summary(), + } + + +def test_ep_prefill_requires_every_local_expert_shard_to_execute(): + summaries = [ + _ep4_summary(rank, hit_count) + for rank, hit_count in enumerate([0, 80, 84, 92]) + ] + + with pytest.raises(RuntimeError, match="did not execute a local routed expert"): + _validate_parallel_evidence( + summaries, + tensor_parallel_size=1, + expert_parallel_size=4, + last_moe_layer=46, + ) + + +def test_ep_single_decode_keeps_zero_hit_as_evidence_without_rejecting_it(): + summaries = [ + _ep4_summary(rank, hit_count) + for rank, hit_count in enumerate([0, 1, 1, 2]) + ] + + evidence = _validate_parallel_evidence( + summaries, + tensor_parallel_size=1, + expert_parallel_size=4, + last_moe_layer=46, + require_all_ep_ranks_hit=False, + ) + + assert evidence["local_expert_hit_counts"] == [0, 1, 1, 2] + assert evidence["all_ep_ranks_hit"] is False + + +def _hybrid_tp4_ep2_summary(world_rank: int) -> dict: + ep_rank = 0 if world_rank < 2 else 1 + expert_ranks = [0, 2] if world_rank % 2 == 0 else [1, 3] + moe_tensor_ranks = [0, 1] if world_rank < 2 else [2, 3] + layer_key = "46" + return { + "world_rank": world_rank, + "parallel": { + "configured": { + "tensor_parallel_size": 4, + "expert_parallel_size": 2, + "data_parallel_size": 1, + "world_size": 4, + }, + "effective": { + "world": { + "rank": world_rank, + "size": 4, + "ranks": [0, 1, 2, 3], + }, + "attention": { + "rank": world_rank, + "size": 4, + "ranks": [0, 1, 2, 3], + }, + "expert": { + "rank": expert_ranks.index(world_rank), + "size": 2, + "ranks": expert_ranks, + }, + "moe_tensor": { + "rank": moe_tensor_ranks.index(world_rank), + "size": 2, + "ranks": moe_tensor_ranks, + }, + "data": { + "rank": 0, + "size": 1, + "ranks": [world_rank], + }, + }, + "attention_replicated_for_ep": False, + }, + "moe_local": { + layer_key: { + "local_expert_start": ep_rank * 32, + "local_expert_end": (ep_rank + 1) * 32, + "local_hit_count": 8, + } + }, + "replica_consistency": { + "last_logits_comparison": "not_applicable_tp_vocab_sharded", + "last_logits_tolerance_ratio": None, + "moe_layers": { + layer_key: { + "topk_ids_mismatch": False, + "topk_weights_tolerance_ratio": 0.0, + "output_tolerance_ratio": 0.0, + } + }, + }, + "decode_cuda_graph": _graph_summary(enabled=True), + } + + +def test_hybrid_tp_ep_validation_checks_groups_and_graph_on_every_rank(): + evidence = _validate_parallel_evidence( + [_hybrid_tp4_ep2_summary(rank) for rank in range(4)], + tensor_parallel_size=4, + expert_parallel_size=2, + last_moe_layer=46, + require_decode_cuda_graph=True, + require_decode_cuda_graph_execution=True, + ) + + assert evidence["hybrid_moe"] is True + assert evidence["local_expert_ranges"] == [ + [0, 32], + [0, 32], + [32, 64], + [32, 64], + ] + assert evidence["expected_groups"]["moe_tensor"] == [[0, 1], [2, 3]] From 5cb85953473e58fa11293ad965da5f29ad68db61 Mon Sep 17 00:00:00 2001 From: CURRENTF <1206013760@qq.com> Date: Fri, 7 Aug 2026 22:02:25 +0800 Subject: [PATCH 17/48] perf: accelerate glm flash runtime --- docs/zh/design/README.md | 2 +- docs/zh/design/glm-4.7-flash-support-plan.md | 1171 ++--------------- .../validate_glm_checkpoint_parallel.py | 26 +- .../flashinfer_custom_all_reduce.py | 192 +++ .../distributed/parallel_context.py | 67 + src/sparsevllm/engine/decode_cuda_graph.py | 21 +- src/sparsevllm/engine/model_runner.py | 16 +- src/sparsevllm/layers/embed_head.py | 2 +- src/sparsevllm/layers/linear.py | 41 +- src/sparsevllm/layers/mla_attention.py | 34 +- src/sparsevllm/models/glm4_moe_lite.py | 300 ++++- src/sparsevllm/models/qwen3.py | 2 + src/sparsevllm/operators/mla_attention.py | 292 +++- src/sparsevllm/operators/moe.py | 71 + src/sparsevllm/operators/sgl_fa3.py | 510 +++++++ src/sparsevllm/operators/sgl_moe.py | 248 ++++ src/sparsevllm/triton_kernel/mla/__init__.py | 4 + .../triton_kernel/mla/decode_schedule.py | 75 ++ src/sparsevllm/triton_kernel/moe.py | 89 +- src/sparsevllm/triton_kernel/moe_config.py | 50 + src/sparsevllm/triton_kernel/silu_and_mul.py | 11 +- tests/test_glm4_moe_lite.py | 71 +- tests/test_glm_cuda_graph.py | 3 + tests/test_mla_attention_layer.py | 63 + tests/test_mla_attention_operator.py | 53 +- tests/test_mla_kernels.py | 30 + tests/test_moe_config.py | 55 + tests/test_operator_providers.py | 62 + tests/test_parallel_context.py | 94 ++ tests/test_qwen3_fp8.py | 1 + tests/test_sgl_fa3.py | 343 +++++ tests/test_sgl_moe.py | 154 +++ tests/test_tp_rpc.py | 4 +- tests/test_triton_moe.py | 49 + .../test_validate_glm_checkpoint_parallel.py | 34 + 35 files changed, 3049 insertions(+), 1191 deletions(-) create mode 100644 src/sparsevllm/distributed/flashinfer_custom_all_reduce.py create mode 100644 src/sparsevllm/operators/sgl_fa3.py create mode 100644 src/sparsevllm/operators/sgl_moe.py create mode 100644 tests/test_sgl_fa3.py create mode 100644 tests/test_sgl_moe.py diff --git a/docs/zh/design/README.md b/docs/zh/design/README.md index 04a9eb1f..b8df59db 100644 --- a/docs/zh/design/README.md +++ b/docs/zh/design/README.md @@ -4,4 +4,4 @@ - [架构](architecture.md) - [Sparse-vLLM 控制图](control-map.md) -- [GLM-4.7-Flash 支持实施计划](glm-4.7-flash-support-plan.md) +- [GLM-4.7-Flash 运行时支持设计](glm-4.7-flash-support-plan.md) diff --git a/docs/zh/design/glm-4.7-flash-support-plan.md b/docs/zh/design/glm-4.7-flash-support-plan.md index 9132c589..0ec711cd 100644 --- a/docs/zh/design/glm-4.7-flash-support-plan.md +++ b/docs/zh/design/glm-4.7-flash-support-plan.md @@ -1,1100 +1,131 @@ -# GLM-4.7-Flash 支持实施计划 +# GLM-4.7-Flash 运行时支持设计 -- 状态：Implemented（H100 BF16；8 种 TP/EP 布局 × 6 种方法的 Graph + Prefix 组合已验证） -- 更新时间：2026-08-05 -- 目标分支：codex/glm-4.7-flash -- 目标模型：zai-org/GLM-4.7-Flash / glm4_moe_lite +本文说明 `glm4_moe_lite` 的稳定运行时契约、组件所有权和支持边界。模型与 +稀疏方法的汇总矩阵见[支持的模型](../features/supported-models.md)。 -## 1. 摘要 +## 设计边界 -本计划为 Sparse-vLLM 增加 GLM-4.7-Flash 的第一阶段支持。实现以 latent MLA 为基础，不通过持久化展开式 K/V 获得临时兼容，也不把 GLM 特例塞进现有显式 K/V 接口。 +GLM-4.7-Flash 使用 BF16 latent MLA。每层、每个 token 的持久化 attention +cache 由 512 维 latent 和 64 维 RoPE key 组成，不持久化展开后的多头 K/V。 -核心决策如下： +实现遵循以下所有权边界： -1. 从固定版本的 LightLLM vendor 最小 MLA Triton 数学核心，移除 LightLLM runtime 依赖，并先以独立 Torch oracle 验证。 -2. 将当前 PrefillComputeView / DecodeComputeView 从固定双 K/V 改为“公共 metadata + tagged payload”。 -3. 保留 CacheManager 对 slot、请求生命周期和 logical view 的所有权；由 storage strategy 管理显式 K/V 或 MLA latent 的物理张量。 -4. 模型只表达 GLM 投影、部分 RoPE、K/V absorption、MoE/router/shared-expert 等语义；模型不选择或直接导入具体 kernel。 -5. MLA kernel 通过 OpSpec -> OpResolver -> Provider -> kernel 路径绑定，初始化后不做运行时静默 fallback。 -6. MoE 只抽取 Qwen3-MoE 中真正通用的 packed-expert 物理执行部分，不构建万能 MoE 模型框架。 -7. GLM response parsing 继续委托 Transformers；本地只提供缺失的声明式 response template。 +- 模型层负责 GLM 投影、部分 RoPE、K/V absorption、Dense/MoE topology、 + biased-sigmoid routing 和 checkpoint 语义。 +- CacheManager 负责 slot、请求生命周期、prefix 生命周期和 logical view。 +- Storage strategy 负责显式 K/V 或 MLA latent 的物理张量、写入和显存核算。 +- MLA operator 通过 `OpSpec -> OpResolver -> Provider -> kernel` 在初始化时 + 绑定；模型不直接选择 kernel。 +- 稀疏方法继续通过 cache-manager-first 接口工作，不在通用 + `attention.py` 中增加 GLM 方法分支。 -本计划最初以 BF16、eager、vanilla latent MLA、TP1 和 TP4/EP1 的短上下文 -正确性为首版目标。后续扩展已在相同 H100/BF16 边界内验证 8 种合法 -`(TP, EP)` 布局，以及 vanilla、StreamingLLM、SnapKV、H2O、OmniKV、 -R-KV 与 decode CUDA Graph、Prefix Cache 的完整交叉组合。当前精确支持边界 -以第 16 节和第 25 节为准；MTP、量化、prefix offload 与实用长上下文仍需 -独立验证。 +Attention compute view 使用公共 metadata 与 tagged payload。显式 K/V 和 +MLA latent payload 是不同类型；消费者收到错误 payload 时必须 fail fast， +不得通过互斥 optional tensor 或 metadata 字典绕过类型契约。 -## 2. 目标与非目标 +## Attention 数据路径 -### 2.1 目标 +Prefill 按以下顺序执行： -- 正确加载 Glm4MoeLiteForCausalLM 的基础模型权重。 -- 正确执行层 0–46；明确识别并跳过 checkpoint 中单独的 MTP 层 47。 -- 使用 latent MLA cache：每层每 token 持久化 512 维 latent 和 64 维 RoPE key。 -- 支持 chunked prefill 的完整历史可见性。 -- 支持 padded static/eager decode，不发生 slot=-1 越界写。 -- 支持 GLM 的首层 Dense、后续 routed MoE、shared expert 和 biased-sigmoid routing。 -- 复用当前 TP/EP、packed expert、provider registry、linear、RMSNorm 和 SwiGLU 基础设施。 -- 提供可复现、artifact-backed 的 tiny 与真实模型验证。 +1. 模型产生当前 chunk 的 latent 和 RoPE key。 +2. Storage 将有效 token 写入持久化 latent cache，并跳过 padding slot。 +3. CacheManager 按 active slot gather 完整可见历史。 +4. MLA layer 临时展开 prefill 所需的 K/V，并调用共享 prefill attention。 +5. 临时 workset 在该次调用结束后释放。 -### 2.2 非目标 +因此，多 chunk prefill 的后续 chunk 必须能看到此前所有可见 token，而不是只 +看到当前 chunk。 -首版不包含： +Decode 使用 absorbed query 直接读取 latent cache。Provider 拥有 Triton +workspace、调度和数值 kernel；输出经过 value projection 重建到模型 hidden +维度。静态 batch 中的 padding row 不得读写 `slot=-1`。 -- MTP/speculative decoding。 -- 非 BF16 checkpoint 或量化权重。 -- CUDA Graph。 -- Prefix Cache 或 prefix offload。 -- SnapKV、PyramidKV、OmniKV、QuEST、H2O、RKV、StreamingLLM、DeltaKV 等稀疏方法。 -- EP 大于 1 的真实模型支持声明。 -- 128K/202K 长上下文吞吐或显存能力声明。 -- LightLLM Python package 或 runtime 依赖。 -- 在 kernel 失败后切换到 Torch/其他 backend 的运行时 fallback。 +## Cache 与前缀复用 -以上列表记录首版冻结范围，不代表后续扩展仍全部未支持。已完成的后续能力和 -仍然保留的拒绝边界见第 16 节兼容矩阵。 +MLA storage 的显存容量按 `512 + 64` 个 BF16 value/token/layer 核算。所有 +allocation、free、reuse、eviction、slot copy 和 prefix replay 都必须同时处理 +latent 与 RoPE cache，且显存统计必须覆盖两个物理张量。 -## 3. 已知模型契约 +Prefix Cache 的模式由稀疏方法决定： -官方配置来源： +- vanilla 和 OmniKV 使用 radix prefix cache。 +- StreamingLLM、SnapKV、H2O 和 R-KV 使用 chain prefix cache。 -- https://huggingface.co/zai-org/GLM-4.7-Flash/blob/main/config.json -- https://github.com/huggingface/transformers/blob/main/src/transformers/models/glm4_moe_lite/modeling_glm4_moe_lite.py +Prefix hit 只有在请求实际复用了 token、cache 状态和方法特定 metadata 时才算 +成功；仅完成请求不能证明 prefix 路径生效。 -关键结构： +## MoE 与并行布局 -| 项目 | 值 | -| --- | ---: | -| model_type | glm4_moe_lite | -| 基础层数 | 47 | -| MTP 层数 | 1 | -| hidden size | 2048 | -| attention heads | 20 | -| KV heads | 20（展开语义） | -| q LoRA rank | 768 | -| kv LoRA rank | 512 | -| qk no-PE dim | 192 | -| qk RoPE dim | 64 | -| qk head dim | 256 | -| value head dim | 256 | -| routed experts | 64 | -| experts per token | 4 | -| shared experts | 1 | -| MoE intermediate size | 1536 | -| Dense intermediate size | 10240 | -| routed scaling factor | 1.8 | -| first Dense layers | 1 | -| dtype | BF16 | +模型第 0 层为 Dense，后续层为 routed MoE，并包含 shared expert。GLM 复用 +Qwen3-MoE 的 packed-expert 物理执行，但保留自己的 router、模型 topology 和 +checkpoint loader。 -必须区分以下维度： +已支持的 `(TP, EP)` 布局为： -~~~ -kv_lora_rank = 512 -rope_dim = 64 -latent_cache_width = 576 -qk_head_dim = 192 + 64 = 256 -softmax_scale = 256 ** -0.5 -~~~ +| TP | EP | Attention | MoE | World size | +| ---: | ---: | --- | --- | ---: | +| 1 | 1 | 单 rank | 单 rank | 1 | +| 2 | 1 | TP=2 | MoE TP=2 | 2 | +| 4 | 1 | TP=4 | MoE TP=4 | 4 | +| 1 | 2 | 每个 EP rank 复制 | EP=2 | 2 | +| 1 | 4 | 每个 EP rank 复制 | EP=4 | 4 | +| 2 | 2 | TP=2 | EP=2，MoE TP=1 | 2 | +| 4 | 2 | TP=4 | EP=2，MoE TP=2 | 4 | +| 4 | 4 | TP=4 | EP=4，MoE TP=1 | 4 | -不能从 latent cache width 576 推导 attention scale，也不能使用 Transformers 映射后的 config.head_dim=64 作为 Sparse-vLLM cache head dimension。 +联合 TP/EP 使用 outer-TP topology：attention TP 为 `T`，MoE EP 为 `E`， +MoE TP 为 `T/E`，world size 为 `T`。该布局要求 `DP=1`、`T % E == 0`， +专家数量能被 `E` 整除，MoE intermediate dimension 能被 `T/E` 整除。 -## 4. 当前运行时约束 +## Operator 与平台边界 -当前实现从抽象层开始假定显式双 K/V： +当前 MLA provider 支持 NVIDIA H100 80GB HBM3、BF16、SM90，以及 TP 1、2、4。 +Provider 在模型初始化时解析并绑定；kernel 执行失败必须直接暴露，不能在 +forward 中静默切换到 Torch 或其他 backend。 -- src/sparsevllm/engine/cache_manager/base.py 的 PrefillComputeView 和 DecodeComputeView 固定包含 k_cache / v_cache。 -- CacheManager 使用 num_key_value_heads 和 head_dim 推导物理 cache shape 与容量。 -- StandardCacheManager 直接分配 K/V tensor，并通过 slot table 暴露计算视图。 -- TritonAttentionBackend 直接消费 view.k_cache / view.v_cache。 -- eager decode 仍可能经过 static batch bucket padding；padding sequence 的 slot_mapping 可以是 -1。 -- chunked prefill 的 view 表示完整可见上下文，而不仅是当前输入 chunk。 +Vendor kernel 固定来源为 LightLLM commit +`65c174ee95ac6a6fd36b18b63d0b33d97e76b770`。本地 vendor 目录保留 +Apache-2.0 license、来源映射和修改说明。模型和 CacheManager 不直接依赖 +LightLLM Python runtime。 -因此禁止以下临时方案： +## Serving 边界 -- 把 latent 和 RoPE tensor 假装成 k_cache / v_cache。 -- 在 ComputeView 中增加一组互斥的 optional tensor 字段。 -- 通过未校验的 metadata 字典把 MLA tensor 旁路给 kernel。 -- 只展开当前 prefill chunk，忽略已缓存历史。 -- 原样复制 LightLLM 的 infer_state wrapper。 +Chat Completions 与 Responses API 都使用 Transformers response parser。本地 +只提供 GLM 缺失的声明式 response template。Terminal EOS、stop boundary、 +流式增量和 raw parser text 的映射由通用 dispatcher/detokenizer 处理，不在 +GLM parser 中加入模型特判。 -## 5. 目标架构 +## 不支持的组合 -~~~ -Glm4MoeLiteForCausalLM - | - +-- Glm4MoeLiteAttention - | | - | +-- q_a/q_b, kv_a/kv_b, partial RoPE, K/V absorption - | | - | +-- MLAAttention (semantic layer) - | | - | +-- AttentionComputeView[AttentionPayload] - | | +-- AttentionViewMeta - | | +-- MlaLatentPayload - | | - | +-- MlaAttentionOpSpec - | +-- OpResolver(DeviceCaps) - | +-- MlaTritonProvider - | +-- vendored LightLLM-derived Triton kernels - | - +-- Glm4MoeLiteMlp - +-- Dense Qwen3MLP (layer 0) - +-- GlmBiasedSigmoidRouter - +-- shared PackedMoeExperts - +-- Qwen3MLP shared expert +以下组合必须在配置或模型构造前明确拒绝，不能落入默认实现： -CacheManager - +-- slot allocator / request lifetime / logical selection - +-- ExplicitKVStorage | MlaLatentStorage -~~~ +- H100 以外的 GPU、非 BF16 checkpoint 和量化权重。 +- `DP>1`、上述矩阵以外的 TP/EP 布局。 +- MTP/speculative decoding；loader 只精确跳过 checkpoint 中的 MTP 层。 +- Prefix offload。 +- PyramidKV、QuEST、SkipKV、DeltaKV，以及多个稀疏方法叠加。 +- 128K/202K 长上下文容量或吞吐支持声明。 + +## 验证门禁 -所有权原则： - -- 模型拥有模型语义和逻辑权重。 -- CacheManager 拥有 slot、请求生命周期和 logical visibility。 -- Storage 拥有物理 cache tensor、store 和内存核算。 -- ComputeView 描述一次 attention 计算可见的 metadata 与 payload。 -- Provider 拥有 kernel-specific layout、workspace、launch config 和调用。 -- Platform / DeviceCaps 拥有设备发现和能力事实。 - -## 6. ComputeView 改造 - -### 6.1 设计 - -建议引入： - -~~~python -@dataclass(frozen=True) -class AttentionViewMeta: - active_slots: torch.Tensor - req_indices: torch.Tensor - context_lens: torch.Tensor - max_context_len: int | None = None - attn_score: torch.Tensor | None = None - temp_slots: torch.Tensor | None = None - - -@dataclass(frozen=True) -class ExplicitKVPayload: - k_cache: torch.Tensor - v_cache: torch.Tensor - backend: str = "dense" - metadata: dict[str, Any] | None = None - - -@dataclass(frozen=True) -class MlaLatentPayload: - latent_cache: torch.Tensor - rope_cache: torch.Tensor - - -AttentionPayload = ExplicitKVPayload | MlaLatentPayload - - -@dataclass -class PrefillComputeView: - meta: AttentionViewMeta - payload: AttentionPayload - - -@dataclass -class DecodeComputeView: - meta: AttentionViewMeta - payload: AttentionPayload -~~~ - -Prefill 和 decode 保留不同类型，因为二者的执行、workspace 和生命周期不同；二者共用 payload 类型和公共 metadata。 - -### 6.2 不变量 - -- 不允许一个 view 同时包含显式 K/V 与 MLA latent。 -- 通用 Attention 必须显式要求 ExplicitKVPayload。 -- MLAAttention 必须显式要求 MlaLatentPayload。 -- payload 类型不匹配时在进入 kernel 前失败。 -- backend 和 backend-specific metadata 只属于 ExplicitKVPayload 或对应 provider，不属于公共 metadata。 -- AttentionViewMeta 不携带模型权重或 kernel workspace。 - -### 6.3 迁移策略 - -先进行不改变行为的机械迁移： - -~~~ -view.k_cache -> view.payload.k_cache -view.v_cache -> view.payload.v_cache -view.active_slots -> view.meta.active_slots -view.context_lens -> view.meta.context_lens -view.backend -> view.payload.backend -~~~ - -所有现有 manager 在该提交中只生成 ExplicitKVPayload。迁移提交不得同时引入 MLA 逻辑，以便将现有测试失败明确归因于接口改造。 - -## 7. Cache storage 契约 - -### 7.1 目标 - -避免创建以下类层次： - -~~~ -GlmStandardCacheManager -GlmSnapKVCacheManager -GlmQuestCacheManager -... -~~~ - -CacheManager 继续按 sparse method 选择；物理 payload 通过 composition 委托给 storage。 - -### 7.2 建议接口 - -~~~python -class AttentionCacheStorage(Protocol): - layout: CacheLayout - - def allocate( - self, - *, - num_layers: int, - num_slots: int, - device: torch.device, - ) -> None: ... - - def layer_payload(self, layer_idx: int) -> AttentionPayload: ... - - def store( - self, - layer_idx: int, - slot_mapping: torch.Tensor, - payload: Any, - ) -> None: ... - - def bytes_per_slot_per_layer(self) -> int: ... - - def accounting_tensors(self) -> tuple[torch.Tensor, ...]: ... -~~~ - -具体类： - -~~~ -ExplicitKVStorage - k_cache [layers, slots, kv_heads, head_dim] - v_cache [layers, slots, kv_heads, head_dim] - -MlaLatentStorage - latent_cache [layers, slots, 1, 512] - rope_cache [layers, slots, 1, 64] -~~~ - -### 7.3 内存核算 - -所有以下逻辑必须从 storage 获取，不再使用全局固定公式 2 * num_kv_heads * head_dim： - -- 每 slot 每 layer 字节数。 -- 最大可分配 slot 数。 -- prefix block 字节数。 -- allocator/debug tensor 汇总。 -- live-token 和 reserved-token 内存报告。 - -accounting_tensors() 必须显式返回 storage 持有的物理 tensor，不能依赖递归扫描任意对象。 - -### 7.4 Padding 安全 - -store() 和 vendored copy kernel 必须跳过 slot_mapping < 0。以下两种模式都必须覆盖： - -- 真实 request slot。 -- static/eager padded batch 的 slot=-1。 - -禁止依赖“关闭 CUDA Graph 后不会 padding”的假设。 - -## 8. LightLLM MLA kernel vendoring - -### 8.1 固定上游 - -- 仓库：https://github.com/ModelTC/lightllm -- commit：65c174ee95ac6a6fd36b18b63d0b33d97e76b770 -- license：Apache-2.0 - -本地目录： - -~~~ -src/sparsevllm/triton_kernel/mla/ -~~~ - -### 8.2 最小来源映射 - -| 上游文件 | 本地计划文件 | 处理方式 | -| --- | --- | --- | -| common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding_stage1.py | decode_stage1.py | 保留 Triton 数学核心，改为显式 tensor/stride/config API | -| common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding_stage2.py | decode_stage2.py | 保留归并数学核心，移除 LightLLM 状态依赖 | -| common/basemodel/triton_kernel/mla_att/decode_att/gqa_flash_decoding.py | decode_schedule.py | 不原样复制；仅参考 stage 调度，本地重写 workspace/launch API | -| common/basemodel/triton_kernel/kv_copy/mla_copy_kv.py | copy_latent.py | 增加 slot=-1 mask 与本地 layout 校验 | -| models/deepseek2/triton_kernel/sample_kv.py | gather_latent.py | 适配 slot-indirected ragged/full-history gather | - -每个衍生文件必须： - -- 保留或补充 Apache-2.0 SPDX 标识。 -- 注明原始仓库、固定 commit 和原文件路径。 -- 注明本地修改。 -- 不 import lightllm。 - -目录同时包含： - -~~~ -README.md -LICENSE.lightllm -~~~ - -### 8.3 明确不复制 - -- LightLLM infer_state。 -- request manager / memory manager。 -- 动态挂载到 infer_state 的 workspace。 -- frozendict。 -- LightLLM global kernel config 或 autotuner。 -- LightLLM device utilities。 -- forward 内的动态 vSM 探测和 allocation。 -- context_flashattention_nopad_with_v.py 作为 GLM 首版 prefill kernel。 - -最后一项不能直接用于 GLM：其契约要求 q_nope_dim == v_dim，而 GLM 是 192 与 256。 - -### 8.4 Vendor kernel API - -Decode 输入： - -~~~ -q_nope_absorbed [batch, local_heads, 512] -q_rope [batch, local_heads, 64] -latent_cache [slots, 1, 512] -rope_cache [slots, 1, 64] -active_slots [rows, max_context_len] -req_indices [batch] -context_lens [batch] -softmax_scale scalar, fixed from qk_head_dim=256 -workspace explicit tensors -launch_config immutable value object -~~~ - -输出： - -~~~ -latent_output [batch, local_heads, 512] -~~~ - -Kernel 不读取 HF config、model type、环境变量或 backend 名称。 - -## 9. MLA operator/provider - -新增： - -~~~ -src/sparsevllm/operators/mla_attention.py -~~~ - -建议 spec： - -~~~python -@dataclass(frozen=True) -class MlaAttentionOpSpec: - num_q_heads: int - kv_lora_rank: int - rope_dim: int - qk_head_dim: int - value_head_dim: int - activation_dtype: torch.dtype - cache_dtype: torch.dtype - tp_size: int - cuda_graph: bool -~~~ - -softmax_scale 可以由经过校验的 qk_head_dim 得出，也可作为只读规范化字段保存；禁止从 kv_lora_rank + rope_dim 得出。 - -Provider 选择链： - -~~~ -MLAAttention semantic call - -> MlaAttentionOpSpec - -> OpResolver(DeviceCaps) - -> MlaTritonProvider - -> prepare workspace/layout - -> decode_stage1/decode_stage2 -~~~ - -约束： - -- provider 在 model construction/ModelRunner 初始化阶段绑定一次。 -- forward hot path 不 resolve provider。 -- supports(spec, caps) 检查 platform、Triton、dtype、维度、TP、graph 和已验证硬件。 -- 未验证的 GPU 架构不宣称支持。 -- 首版只有一个 production provider；Torch 只用于测试 oracle。 -- kernel/JIT 失败直接暴露，不 catch 后切换实现。 -- Provider 使用 Platform / DeviceCaps，通用代码不直接探测 torch.cuda。 - -LightLLM 固定 commit 没有目标 H100/GLM shape 的现成 tuning config，因此首版使用明确、安全的静态 launch config；只有 H100 数值和性能实测后才将该 shape 标记为已验证。 - -## 10. Prefill 数据路径 - -MLA persistent cache 只保存 latent/rope，但 chunked prefill 必须看到完整历史。 - -正确顺序： - -1. 当前 chunk 计算 q_a/q_b 和 kv_a。 -2. 对当前 chunk 的 512 维 latent 做 RMSNorm。 -3. 对当前 chunk 的 64 维 key 应用 RoPE。 -4. 通过 padding-safe store 写入持久 latent/rope cache。 -5. 根据 active_slots + req_indices + context_lens gather 每条序列的完整历史 latent/rope，包括之前 chunk 和当前 chunk。 -6. 使用 kv_b_proj 将完整历史 latent 临时展开为： - - k_nope [visible_tokens, local_heads, 192] - - v [visible_tokens, local_heads, 256] -7. 将历史 RoPE key 扩展/映射到 local heads，与 k_nope 形成 256 维 K。 -8. 使用 Sparse-vLLM 现有 contiguous/slot-compatible 256 维 prefill attention kernel。 -9. prefill 完成后释放 gather/expanded workspace，不将其写回持久 cache。 - -需要单独的临时 workset： - -~~~python -@dataclass -class MlaPrefillWorkset: - gathered_latent: torch.Tensor - gathered_rope: torch.Tensor - expanded_k: torch.Tensor - expanded_v: torch.Tensor - packed_offsets: torch.Tensor -~~~ - -Workspace 必须有显式显存预算和有界分配。无法满足预算时应在运行 kernel 前报错，不允许降级为缺失历史或缩短上下文。 - -## 11. Decode 数据路径 - -1. 当前 token 计算： - - q_nope [batch, local_heads, 192] - - q_rope [batch, local_heads, 64] - - latent [batch, 1, 512] - - k_rope [batch, 1, 64] -2. padding-safe store 将 latent/k_rope 写入当前 slot。 -3. 使用 K_b 将 query no-PE 部分吸收到 latent 空间： - - ~~~ - q_nope_absorbed [batch, local_heads, 512] - ~~~ - -4. MLA provider 直接对 MlaLatentPayload 计算 attention，输出： - - ~~~ - latent_output [batch, local_heads, 512] - ~~~ - -5. 使用 V_b 将输出恢复为： - - ~~~ - value_output [batch, local_heads, 256] - ~~~ - -6. flatten local heads，进入 row-parallel o_proj。 - -Decode 计算中的 logits 是： - -~~~ -q_nope_absorbed @ latent_cache.T -+ q_rope @ rope_cache.T -~~~ - -整体使用 256 ** -0.5 scale。 - -## 12. 模型实现 - -新增： - -~~~ -src/sparsevllm/models/glm4_moe_lite.py -~~~ - -模型文件负责： - -- token embedding / final norm / LM head。 -- q_a_proj、q_a_layernorm、q_b_proj。 -- kv_a_proj_with_mqa、kv_a_layernorm、kv_b_proj。 -- 192+64 partial/interleaved RoPE。 -- K/V absorption 和 V reconstruction。 -- Dense/MoE layer topology。 -- checkpoint 逻辑名称映射。 - -模型文件不负责： - -- import vendored kernel。 -- 选择 provider。 -- 探测 GPU。 -- 分配持久 cache。 -- 维护 slot table。 -- 运行时 backend fallback。 - -TP 布局建议： - -- q_a_proj、kv_a_proj_with_mqa 和对应 low-rank norm：replicated。 -- q_b_proj、kv_b_proj：按 local attention heads 做 column/local-head shard。 -- o_proj：row parallel。 -- 每 rank 持有本 rank 的 K_b/V_b head slices。 -- persistent latent/rope cache 在各 TP rank 上语义一致。 - -首版门禁以 TP1 和 TP4 为目标；本轮同时补测 TP2。支持声明只列入真实 -checkpoint 已实测的 `TP=1/2/4`，不从配置可构造性外推其他 TP 规模。 - -## 13. MoE 复用边界 - -### 13.1 抽取内容 - -从 src/sparsevllm/models/qwen3_moe.py 抽取 packed-expert 物理执行： - -- W13/W2 storage。 -- unquantized TP intermediate shard。 -- EP local expert ownership。 -- MoeOpSpec 和 provider resolution。 -- projection loading。 -- local expert completeness validation。 - -建议形成共享 PackedMoeExperts，由 Qwen3-MoE 与 GLM 分别组合。 - -### 13.2 不抽取内容 - -首轮不抽取： - -- Qwen softmax router。 -- 整个 Qwen MoE block。 -- model-specific chunking/all-reduce/debug。 -- checkpoint regex。 -- remote-expert skip bookkeeping。 -- MiniMax packed-expert 实现。 - -这些内容保留在模型层，以避免为了两个不同模型构建过度通用的 MoE framework。 - -### 13.3 GLM MoE - -~~~ -layer 0: - Qwen3MLP(hidden=2048, intermediate=10240) - -layer 1-46: - routed = PackedMoeExperts(x, topk_ids, topk_weights) - shared = Qwen3MLP(hidden=2048, intermediate=1536)(x) - output = world_reduce(routed) + shared -~~~ - -Router 必须复现： - -1. FP32 gate linear。 -2. sigmoid。 -3. correction bias 只影响 expert selection。 -4. 从原始 sigmoid score gather selected weights。 -5. Top-4 normalization。 -6. 乘 routed_scaling_factor=1.8。 - -当前 MiniMax router kernel 硬编码 256 experts/Top-8。将其泛化为支持 GLM 64/Top-4 应作为独立提交；不复制一个 GLM 专属基础 kernel，也不在 packed-expert 抽取提交中顺带迁移 MiniMax storage。 - -## 14. 权重加载与 MTP 边界 - -Loader 要求： - -- 所有基础层 0–46 权重必须被消费。 -- 所有本地 EP experts 的 gate/up/down 权重必须完整。 -- 只允许跳过其他 EP rank 持有的 expert 权重。 -- e_score_correction_bias 保持 FP32。 -- 只允许精确跳过 model.layers.47.* MTP 权重。 -- 任何其他 missing/unexpected tensor 直接失败。 - -加载结束时记录： - -- loaded tensor count。 -- local expert shard count。 -- remote EP skipped count。 -- MTP skipped count。 -- 其他 skipped/missing/unexpected count。 - -MTP 不得通过宽泛 silent ignore 机制掩盖错误层数或命名变化。 - -## 15. Response parsing - -Sparse-vLLM 继续使用 Transformers parser： - -~~~ -tokenizer.parse_response() # 非流式 -tokenizer.get_response_parser() # 流式 -~~~ - -官方 GLM tokenizer 当前缺少 response_template，因此本地仅增加声明式 fallback，描述： - -- ...。 -- function-name ... 。 -- ...。 -- ...。 - -禁止实现另一套 XML parser、streaming state machine 或 tool-call parser。 - -CLI 建议新增通用 auto；现有 qwen3、minimax_m2 保留兼容，glm47 仅作为 alias，不选择另一套 parser 实现。 - -## 16. 当前兼容矩阵 - -| 能力 | 当前状态 | 门禁 | -| --- | --- | --- | -| BF16 | 已实现并验证 | tiny + 真实权重验证 | -| NVIDIA H100 80GB HBM3 | 已实现并验证 | `triton_h100` provider 实机门禁 | -| eager | 已实现并验证 | `enforce_eager=True` | -| vanilla latent MLA | 已实现并验证 | kernel/storage/model 全链路 | -| TP1/EP1 | 已实现并验证 | 真实模型 | -| TP2/EP1 | 已实现并验证 | 真实模型 32 步 greedy 对齐 | -| TP4/EP1 | 已实现并验证 | 真实模型、multi-chunk、ragged batch 与 API 并发 | -| TP1/EP2、TP1/EP4 | 已实现并验证 | replicated attention + EP，跨 rank router/MoE/logits 一致性与 expert 覆盖 | -| TP2/EP2、TP4/EP2、TP4/EP4 | 已实现并验证 | outer TP：attention TP=`T`、MoE EP=`E`、MoE TP=`T/E`、world=`T` | -| decode CUDA Graph | 8 种布局均支持 6 种方法 | 每个 world rank 的 capture/replay、零 eager/force-eager fallback | -| Prefix Cache | 8 种布局均支持 6 种方法 | vanilla/OmniKV radix；StreamingLLM/SnapKV/H2O/R-KV chain | -| CUDA Graph + Prefix Cache | 8 种布局 × 6 种方法已联合验证 | 48/48 真实 checkpoint 组合门禁 | -| sparse methods | 8 种布局均支持 StreamingLLM、SnapKV、H2O、OmniKV、R-KV | H2O 维持实验性标记；TP>1 使用 TP-local selection | -| Prefix offload | 暂不支持 | 配置阶段拒绝 | -| PyramidKV、QuEST、SkipKV、DeltaKV | 暂不支持 | method registry 显式拒绝 | -| MTP | 暂不支持 | loader 精确跳过 | -| quantization | 暂不支持 | checkpoint config 拒绝 | -| 128K/202K | 未验证 | 不写入支持声明 | - -配置校验必须发生在模型构造和 cache 分配之前，给出具体不兼容原因。 - -## 17. 分阶段执行计划 - -### 阶段 0：同步与基线 - -工作： - -- 同步当前分支与目标 main 基线。 -- 确认 worktree 干净。 -- 记录 Python、Torch、Transformers、Triton、CUDA 和 GPU 信息。 -- 校验本地 checkpoint 是否具备 48 个 shard、index 和可追踪 revision/hash。 -- 将 Hugging Face、Xet、Triton、TorchInductor、临时文件和实验输出路由到 - 非系统盘 `${SPARSEVLLM_ARTIFACT_ROOT}`。 - -门禁： - -- 不覆盖现有 partial checkpoint。 -- checkpoint 不完整时不运行真实模型加载。 -- GPU 任务前重新检查全部设备，仅使用空闲设备。 - -### 阶段 1：Vendor 独立 MLA kernels - -首个提交： - -~~~ -feat: vendor lightllm mla kernels -~~~ - -内容： - -~~~ -src/sparsevllm/triton_kernel/mla/ - __init__.py - decode_stage1.py - decode_stage2.py - decode_schedule.py - copy_latent.py - gather_latent.py - README.md - LICENSE.lightllm -tests/test_mla_kernels.py -~~~ - -门禁： - -- 零 LightLLM runtime import。 -- 所有函数使用显式 tensor/stride/workspace/config API。 -- copy 正确跳过 slot=-1。 -- gather 支持 ragged、non-contiguous、完整历史和 padded rows。 -- Torch oracle 覆盖 local heads 5/10/20 和边界长度。 -- 空闲 H100 数值 smoke 通过后才合入。 - -该阶段不接 Model、CacheManager 或 OpenAI serving。 - -### 阶段 2：ComputeView tagged payload 迁移 - -工作： - -- 引入 AttentionViewMeta、ExplicitKVPayload、MlaLatentPayload。 -- 将现有 prefill/decode view 改为 meta + payload。 -- 迁移所有现有 manager/backend/tests 到 ExplicitKVPayload。 - -门禁： - -- 不加入 MLA 执行逻辑。 -- 当前 explicit-KV 测试全部保持通过。 -- git diff 中不存在 optional payload field bag。 -- 通用 backend 在 payload 错型时 fail fast。 - -### 阶段 3：Storage composition - -工作： - -- 引入 ExplicitKVStorage 和 MlaLatentStorage。 -- StandardCacheManager 委托 allocation、store、layer payload、bytes 和 accounting。 -- 现有 sparse managers 首版只接受 explicit layout。 - -门禁： - -- explicit storage 与迁移前分配量和结果一致。 -- MLA storage 容量按 576 BF16 values/token/layer 计算。 -- slot reuse/free/padding 不泄漏旧数据。 -- 所有不支持组合在配置阶段拒绝。 - -### 阶段 4：MLA operator/provider 与 decode - -工作： - -- 新增 MlaAttentionOpSpec、registry 和 Triton provider。 -- 初始化时绑定 provider。 -- 构造显式 workspace 和安全 launch config。 -- 打通 absorbed query -> latent decode -> V_b reconstruction。 - -门禁： - -- resolver supported/rejected tests。 -- 不在 forward 中 resolve 或动态选择 provider。 -- softmax scale 使用 256，而不是 576。 -- decode 与 Torch/Transformers component oracle 对齐。 -- static padded batch 不越界。 - -### 阶段 5：完整历史 chunked prefill - -工作： - -- 保存当前 chunk latent/rope。 -- gather 完整 visible history。 -- 临时展开 K/V。 -- 接现有 256 维 prefill attention kernel。 -- 释放 workset。 - -门禁： - -- 单 chunk 与多 chunk 输出对齐。 -- 第二及后续 chunk 能看到历史 token。 -- ragged batch、prefix length 和不同 chunk boundary 对齐。 -- workspace 有界且在不足时明确失败。 - -### 阶段 6：Tiny MLA/model contract - -工作： - -- 扩展 tiny-random config/checkpoint generator。 -- 建立两层模型：第 0 层 Dense、第 1 层 MoE。 -- 与 Transformers 比较 projection、RoPE、router、attention、MLP 和 logits。 - -门禁： - -- prefill logits 对齐。 -- 多步 decode greedy token 对齐。 -- TP rank-local projection slice 正确。 -- 所有 logits finite。 - -### 阶段 7：最小 packed-expert 抽取 - -工作： - -- 抽取 Qwen3-MoE packed expert 物理执行。 -- Qwen3-MoE 使用共享实现，保持语义不变。 -- GLM 准备复用相同 packed experts。 - -门禁： - -- Qwen3-MoE 现有 config、loader、provider、TP/EP 测试全部通过。 -- 不抽取 router、model block、checkpoint regex 或 MiniMax storage。 - -### 阶段 8：GLM MoE、模型与 loader - -工作： - -- 泛化 biased-sigmoid router 到 64/Top-4。 -- 实现 Dense/MoE/shared-expert topology。 -- 注册 model runner/config/runtime topology。 -- 实现 strict loader 和 MTP skip。 - -门禁： - -- tiny end-to-end 对齐。 -- loaded/skipped/missing tensor 统计满足精确约束。 -- Qwen3-MoE 和 MiniMax router 回归测试通过。 - -### 阶段 9：Transformers response template - -工作： - -- 增加 GLM declarative template fallback。 -- 加入非流式与流式 OpenAI 响应测试。 - -门禁： - -- 实际调用 Transformers parse_response / get_response_parser。 -- thinking、普通 content、多 tool calls、Unicode、跨 chunk tag、畸形输出均覆盖。 - -### 阶段 10：真实模型验证 - -顺序： - -1. checkpoint 完整性与 loader metadata 检查。 -2. HF reference 与 Sparse-vLLM 顺序执行，避免同时占用显存。 -3. TP1、eager、vanilla、短上下文 smoke。 -4. TP1 component/logit/token parity。 -5. 等待四张 GPU 同时空闲。 -6. TP4/EP1 smoke。 -7. TP1/TP4 greedy token 与 logits 对齐。 -8. OpenAI chat/reasoning/tool call smoke。 - -首轮真实模型建议： - -~~~ -max_model_len: 2048-4096 -prompt: 128-512 tokens -decode: 32-64 tokens -temperature: 0 -dtype: bfloat16 -decode_cuda_graph: false -enable_prefix_caching: false -vllm_sparse_method: "" -~~~ - -## 18. 测试矩阵 - -### 18.1 Kernel tests - -- batch：1、2、8。 -- local heads：5、10、20。 -- context：1、31、32、33、127、128、129、255、256、257、1024、4096。 -- slot：contiguous、ragged、non-contiguous、invalid duplicate rejection、padding -1。 -- dtype：BF16；其他 dtype 明确拒绝或单独验证。 -- stage1/stage2 分别与整体 oracle 对齐。 -- workspace 最小容量和不足容量失败。 - -### 18.2 ComputeView/storage tests - -- payload 类型不匹配失败。 -- explicit path 行为不变。 -- bytes-per-slot 与实际 tensor bytes 一致。 -- allocate/free/reuse 不泄漏。 -- accounting 覆盖所有 storage tensor。 -- 不支持的 sparse/graph/prefix 配置构造前失败。 - -### 18.3 Component parity - -- q_a/q_b projection。 -- kv_a norm 和 kv_b split。 -- 192+64 interleaved RoPE。 -- prefill expanded K/V。 -- decode K_b absorption 和 V_b reconstruction。 -- biased-sigmoid ids/weights/scaling。 -- routed + shared expert 输出。 -- Dense/MoE layer topology。 - -每个测试使用显式 atol/rtol。BF16 kernel 初始可用 3e-2 量级作为探测值，但最终容差必须由误差分布和独立 reference 决定，不能通过不断放宽容差让测试通过。 - -### 18.4 End-to-end parity - -- tiny prompt prefill logits。 -- tiny 多步 greedy decode。 -- 真实模型首 token logits/top-k margin。 -- 真实模型短序列 greedy token。 -- TP1 与 TP4 对齐。 -- TP1/EP2 与 TP1/EP4 的跨 rank router、MoE 输出和 logits 一致性。 -- 8 种合法 `(TP, EP)` 布局的 group topology、expert 分片与执行覆盖。 -- vanilla/OmniKV radix 与 StreamingLLM/SnapKV/H2O/R-KV chain prefix cache - 的真实 hit 和 token reuse。 -- vanilla、StreamingLLM、SnapKV、H2O、OmniKV、R-KV 的真实模型 - static/eager 与 decode CUDA Graph 全词表 logits 对比；同时要求 Graph capture、 - business-request replay、零 eager/force-eager fallback 和方法特定触发证据。 -- 以上 6 种方法与 8 种并行布局、decode CUDA Graph、Prefix Cache 的完整 - 交叉组合；不能用单特性分别通过替代联合生命周期门禁。 - -非空输出不构成正确性证明。 - -## 19. 可观测性与实验产物 - -每次 GPU 验证至少保存： - -- code commit 和 dirty status。 -- 完整命令与 resolved config。 -- Python/Torch/Transformers/Triton/CUDA 版本。 -- GPU 型号和 device index。 -- model/checkpoint revision/hash/completeness。 -- selected MLA provider 与 rejection summary。 -- cache layout、每 token 字节数、实际 slot capacity。 -- loaded/skipped/missing weight counts。 -- raw outputs、parsed outputs、per-sample status、aggregate result。 -- stdout/stderr 和错误扫描。 - -建议通过环境变量配置非系统盘输出根目录： - -~~~ -${SPARSEVLLM_ARTIFACT_ROOT}/glm-4.7-flash// -~~~ - -系统盘空间不足时不得把 Hugging Face、Triton、TorchInductor 或实验临时文件写入默认 home cache。 - -## 20. 风险登记 - -| 严重度 | 风险 | 最小修正 | -| --- | --- | --- | -| P0 | chunked prefill 只展开当前 chunk，历史不可见 | 先 store，再按 active slots gather 完整历史，然后展开 | -| P0 | LightLLM copy 对 slot=-1 无 mask | vendor 时增加 padding mask 和回归测试 | -| P0 | 用 576 推导 softmax scale | OpSpec 独立记录 qk head dim 256 | -| P1 | latent 冒充 k_cache/v_cache | ComputeView tagged payload | -| P1 | model-specific CacheManager 形成组合爆炸 | CacheManager + storage composition | -| P1 | 原样复制 LightLLM wrapper | 只 vendor 数学核心，本地重写 provider/schedule | -| P1 | runtime kernel failure 后静默 fallback | 仅初始化时 resolver fallback；执行失败直接暴露 | -| P1 | 不支持组合落入 dense/default validation | GLM compatibility 在配置阶段显式拒绝 | -| P2 | prefill workspace 随历史无界增长 | 显式预算、容量估算和 fail-fast | -| P2 | H100 没有固定上游 tuning config | 安全静态 config + 实机数值/性能验证 | -| P2 | MoE 抽取范围过大 | 只抽 packed expert 物理执行 | -| P3 | vendor 来源/修改不清 | README、固定 commit、SPDX、license 和 source mapping | - -## 21. 预计文件影响 - -| 区域 | 文件 | -| --- | --- | -| ComputeView | src/sparsevllm/engine/cache_manager/base.py、所有 view 构造/消费点 | -| Storage | src/sparsevllm/engine/cache_manager/storage/*、standard.py | -| Attention layer | src/sparsevllm/layers/mla_attention.py、attention.py | -| Operator | src/sparsevllm/operators/mla_attention.py、registry tests | -| Kernels | src/sparsevllm/triton_kernel/mla/* | -| Model | src/sparsevllm/models/glm4_moe_lite.py | -| MoE | qwen3_moe.py、共享 packed-expert 模块、router kernel | -| Config/runtime | configs/model.py、configs/runtime.py、method_registry.py | -| Dispatch/warmup | engine/model_runner.py、engine/llm_engine.py | -| Cache factory | engine/cache_manager/base.py 或后续独立 factory | -| Serving | entrypoints/openai/serving/response_parsing.py、api_server.py | -| Tiny/tests | debug/tiny_random.py、新增 kernel/model/parser tests | - -## 22. 提交拆分 - -计划提交： - -1. feat: vendor lightllm mla kernels -2. refactor: tag attention payload views -3. feat: add mla cache storage -4. feat: add mla attention provider -5. feat: add mla prefill and decode -6. refactor: share packed moe experts -7. feat: add glm4 moe lite model -8. feat: add glm response template -9. test: validate glm4 runtime parity -10. docs: document glm4 support matrix - -每个提交要求： - -- git diff --check 通过。 -- 对应的 targeted tests 通过。 -- 不混入无关重构。 -- 不提交失败或不完整的真实实验结果作为成功结论。 - -ComputeView 迁移与 MLA 引入分开；Qwen MoE 抽取与 GLM model 引入分开；支持矩阵只在真实门禁完成后更新。 - -## 23. 完成定义 - -只有以下条件全部满足，才可标记“GLM-4.7-Flash 基础支持完成”： - -- Vendor kernel 有固定来源、license、修改说明和独立 oracle tests。 -- ComputeView 不再假定唯一 payload 是显式双 K/V。 -- 所有原有 explicit-KV 测试保持通过。 -- Persistent cache 只保存 512+64 latent/rope，不保存展开式 20-head K/V。 -- Chunked prefill 包含完整可见历史。 -- Padded decode 不对 slot=-1 写入。 -- MLA scale 使用 256 ** -0.5。 -- Provider 初始化时绑定，运行时无静默 fallback。 -- Tiny model 的 attention、MoE、logits 和多步 decode 与 Transformers 对齐。 -- 真实 checkpoint 除精确 MTP/远端 EP allowlist 外无 missing/unexpected weights。 -- TP1/TP2/TP4、EP1 真实短序列验证通过；其他 TP/EP 规模不得外推。 -- OpenAI 非流式/流式 reasoning 和 tool-call parsing 通过 Transformers parser。 -- Artifact 能证明实际执行了目标 MLA provider、latent cache 和 GLM MoE 路径。 -- 文档支持矩阵准确区分已实现、已验证和暂不支持。 - -## 24. 开始执行前检查清单 - -- [ ] 当前分支同步到目标 main。 -- [ ] worktree 干净或已有改动已明确归属。 -- [ ] LightLLM source commit 固定。 -- [ ] vendor source/license manifest 准备完成。 -- [ ] GPU 状态重新检查。 -- [ ] 系统盘 cache/output 路径迁移到非系统盘 `${SPARSEVLLM_ARTIFACT_ROOT}`。 -- [ ] 首个 kernel commit 的 Torch oracle 写好。 -- [ ] ComputeView 迁移单独排期，不与 MLA 功能混合。 -- [ ] checkpoint 完整性门禁准备好。 -- [ ] 所有首版不支持组合有明确 config rejection。 - -第一步是独立完成并验证 feat: vendor lightllm mla kernels。该提交通过前，不进入 ComputeView、CacheManager 或模型实现。 - -## 25. 当前支持状态与验证摘要 - -当前里程碑已经完成 BF16 latent MLA，以及 `(TP, EP)` 为 `(1,1)`、`(2,1)`、 -`(4,1)`、`(1,2)`、`(1,4)`、`(2,2)`、`(4,2)`、`(4,4)` 时，vanilla、 -StreamingLLM、SnapKV、H2O、OmniKV、R-KV 与 decode CUDA Graph、Prefix -Cache 的完整交叉组合，硬件限定为 NVIDIA H100 80GB HBM3、`DP=1`。 -支持声明只覆盖这些实测组合，不向其他并行规模、其他 sparse 方法、prefix -offload 或多稀疏算法叠加外推。 - -### 25.1 已落地的架构边界 - -- MLA kernel 固定来源为 LightLLM commit - `65c174ee95ac6a6fd36b18b63d0b33d97e76b770`，本地保留 Apache-2.0 - license、逐文件来源映射和修改说明。 -- ComputeView 使用 tagged attention payload；explicit K/V 与 MLA latent payload - 通过类型契约区分，不使用互斥 optional tensor 或 metadata 字典旁路。 -- CacheManager 继续拥有 slot、生命周期和 logical view；storage strategy 分别管理 - explicit K/V 与每层每 token `512 + 64` BF16 latent/rope cache。 -- MLA provider 通过 operator resolver 在初始化时绑定；实机选择 - `triton_h100`，执行失败不会静默切换实现。 -- GLM MoE 复用 Qwen3-MoE 的 packed expert 物理执行，只保留 GLM 自己的 - biased-sigmoid router、Dense/MoE topology 和 checkpoint 语义。 -- 联合 TP/EP 复用 Qwen3-MoE 的 outer-TP group topology：attention TP 为 - `T`、MoE EP 为 `E`、MoE TP 为 `T/E`、world size 为 `T`；纯 EP 保留 - replicated-attention world size `E`，纯 TP 保留普通 tensor parallel。 -- Chat Completions 与 Responses API 均继续调用 Transformers response parser； - 本地只提供 GLM 声明式 template。Terminal EOS 和 stop boundary 在 dispatcher / - detokenizer 通用边界处理，不在 GLM parser 中写模型特判。 -- 稀疏方法继续复用 cache-manager-first 路径；MLA storage 提供 latent cache 的 - eviction、query-key materialization 和 graph-stable selection metadata，不在 - `attention.py` 增加 GLM 方法分支。 -- Decode CUDA Graph 保留 warmup graph 及其 private-pool ownership；业务请求用 - capture/replay counter 增量验证实际 replay，不通过销毁 warmup graph 后重捕获来 - 重置测量。 - -### 25.2 验证结果 - -| 门禁 | 结果 | 证据 | -| --- | --- | --- | -| CPU regression | `1523 passed, 172 skipped, 239 subtests passed` | CPU-only 完整回归 | -| CUDA kernel/operator/model targeted regression | `83 passed` | H100 kernel、operator、attention layer、模型与稀疏 score lifecycle targeted selection | -| Tiny multi-chunk TP1 | 17-token prompt、chunk size 8，实际 3 chunks；prefill 与两步 decode argmax 全匹配，max abs diff 分别为 `0.00390625`、`0.00390625`、`0.0048828125` | run manifest、raw comparison 与 stdout | -| 真实 checkpoint TP1 | 48/48 shards、9491 tensors、55.77 GiB、8832 local expert shards、精确跳过 212 个 MTP tensors；prefill + 32 decode steps argmax 全匹配 | run manifest、raw comparison 与 stdout | -| 真实 checkpoint TP2 | 每 rank 28.01 GiB，两个 rank 均加载 48/48 shards、9491 tensors；prefill + 32 decode steps 共 33/33 argmax 全匹配，greedy token 序列与 TP1 相同 | run manifest、raw comparison 与 stdout | -| 真实 checkpoint TP4/EP1 | 每 rank 14.13 GiB，四个 rank 均绑定 `triton_h100` MLA 与 Triton MoE；短序列 33/33、257-token 五段 prefill + decode 9/9 argmax 全匹配 | run manifest、raw comparison 与 stdout | -| TP4 synthetic API ragged + natural LCC | 64/256/768/1536-token synthetic API batch 连续三轮共 12/12 成功且跨轮逐字稳定；另以 1004/1023/1024/1025-token natural LCC 样本对 HF teacher-forced prefill 7/7、decode 27/28，唯一 teacher-forced argmax mismatch 为两个 HF 候选 logit 完全相等的 BF16 tie | raw/parsed responses、per-row comparison 与 request logs | -| TP4 OpenAI client gates | Chat 普通/思考、Chat SSE、Responses、Responses SSE、function tool call、streaming tool call 和 4-way concurrency 共 8/8 client gates 通过；parser 实际委托 Transformers。服务在验证后由 operator 主动停止，teardown 不作为 client gate | command/env manifest、raw parser events、parsed responses、validation status 与 request logs | -| 真实 checkpoint EP2/EP4 | `TP=1, DP=1` 下 EP2 与 EP4 均成功；各 rank 的 router top-k、MoE 输出和最终 logits 精确一致，prefill 覆盖全部 expert shard | summary、完整 logits 与跨 rank comparison | -| 真实 checkpoint 联合 TP/EP topology | TP2/EP2、TP4/EP2、TP4/EP4 均成功；attention、expert、MoE-TP、data group 与 outer-TP 契约逐 rank 一致，所有 expert partition 在 prefill 实际命中，router top-k 与 MoE 输出跨 rank 精确一致 | per-rank group summary、expert hit count、router/MoE comparison 与 status | -| TP2/EP1 与 TP2/EP2 数值门禁 | 同 prompt/config 的 prefill logits max/mean/p99 abs diff 为 `1.5/0.174635/0.5`，final decode 为 `2.0/0.417368/1.0`；首 token 差异来自 reference 中两个候选的零 margin BF16 tie，后续 token 完全一致 | 两侧 float32 全词表 logits、token 序列、SHA256 与独立 comparison status | -| 真实 checkpoint Prefix Cache | vanilla/OmniKV radix 与 StreamingLLM/SnapKV/H2O/R-KV chain 共 6 种模式成功；每种均完成 2/2 请求，第二请求命中，eligible hit rate 为 1.0，chain 模式无 invalidation/prefix mismatch | aggregate metrics、resolved config 与 request logs | -| 真实 checkpoint decode CUDA Graph | vanilla、StreamingLLM、SnapKV、H2O、OmniKV、R-KV 共 6 种方法均通过 10/10 gates；两段 prompt 产生的 6 行全词表 logits shape 为 `[6, 154880]`，eager/Graph max abs diff 均为 `0.0`；每项 lifetime capture 为 1、业务 replay 增量为 6，且无 eager/force-eager fallback，方法特定 trigger 均有证据 | 完整 logits、runtime counter before/after/delta、selection/cache trace 与 status | -| Graph + Prefix 完整组合矩阵 | 8 种 `(TP, EP)` 布局 × 6 种方法共 48/48 case 成功；每项完成 2/2 请求且 eligible hit rate 为 1.0，每个 world rank 均 lifetime capture=1、业务 replay 增量=4、eager/force-eager 增量=0 | per-case metrics、per-rank Graph before/after/delta、prefix/chain lifecycle 与 aggregate status | - -完整命令、依赖与设备版本、dirty status、raw/parsed output 和本机 artifact -路径保存在私有实验记录中，不写入公开仓库文档；上表只保留稳定的支持边界和 -可复现门禁摘要。 - -TP1、TP2 和 TP4 的真实 checkpoint 32 步 greedy token 序列全部一致，足以证明 -该短序列的独立 greedy rollout 一致。它不等价于严格逐元素 logits parity： -后续步骤及 multi-chunk case 的 BF16 raw logit max abs diff 可到 `9.0`，并带 -近似全局平移；自然文本 teacher-forced 测试还观测到一次零 margin tie。当前证据 -只支持“非 tie 的 argmax 与短序列 greedy 行为一致”，不支持“所有 logits 数值 -严格等价”。 - -### 25.3 尚未完成的门禁 - -- Prefix Cache offload。 -- PyramidKV、QuEST、SkipKV、DeltaKV 等未列入当前 GLM 矩阵的方法。 -- 量化 checkpoint 与 MTP/speculative decoding。 -- 128K/202K 长上下文容量与吞吐验证。 -- `TP=5` 等其他可整除配置与 `DP>1`。 -- SnapKV + H2O 等多稀疏算法叠加；当前“组合支持”指一个稀疏方法与 - TP/EP、CUDA Graph、Prefix Cache 的能力交叉，不定义新的混合稀疏算法。 -- H100 以外的 GPU 架构。 - -这些组合在完成独立验证前必须继续 fail fast 或维持未支持声明。 +支持声明至少需要覆盖以下可复现门禁： + +- Kernel：BF16 数值 oracle、ragged/non-contiguous slot、边界长度、padding 和 + workspace 容量检查，入口见 `tests/test_mla_kernels.py`。 +- Operator/layer：resolver 拒绝原因、初始化时绑定、prefill 完整历史和 decode + 数值契约，入口见 `tests/test_mla_attention_operator.py` 与 + `tests/test_mla_attention_layer.py`。 +- Storage/lifecycle：allocation、free、reuse、copy、eviction、prefix replay 和 + 显存核算，入口见 `tests/test_attention_cache_storage.py` 与 + `tests/test_glm_mla_prefix_cache.py`。 +- Model/MoE：projection、RoPE、router、packed experts、loader allowlist 和 + tiny 多步 decode，入口见 `tests/test_glm4_moe_lite.py`。 +- CUDA Graph：真实 capture/replay counter、零 eager fallback、全词表 logits 和 + 方法特定状态证据，入口见 `tests/test_glm_cuda_graph.py` 与 + `tests/test_glm_mla_sparse_methods.py`。 +- Serving：非流式、SSE、reasoning、tool call、EOS 和 stop boundary，入口见 + `tests/test_openai_api_server.py`。 + +验证产物应分别保存 resolved config、命令、代码 revision/dirty status、环境与 +checkpoint 信息、raw/parsed output、逐样本状态和聚合结果。公开文档只维护稳定 +契约、支持边界与自动化门禁，不记录单次运行的通过数量、设备占用或本地路径。 diff --git a/scripts/validation/validate_glm_checkpoint_parallel.py b/scripts/validation/validate_glm_checkpoint_parallel.py index a94df294..c2b5b5f2 100644 --- a/scripts/validation/validate_glm_checkpoint_parallel.py +++ b/scripts/validation/validate_glm_checkpoint_parallel.py @@ -23,6 +23,9 @@ import torch +REPO_ROOT = Path(__file__).resolve().parents[2] + + def _parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Validate one real GLM checkpoint TP/EP variant." @@ -66,13 +69,22 @@ def _write_json(path: Path, payload: Any) -> None: def _git_output(*args: str) -> str: - result = subprocess.run( - ["git", *args], - check=False, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - text=True, - ) + command = ["git", *args] + try: + result = subprocess.run( + command, + cwd=REPO_ROOT, + check=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + except (FileNotFoundError, subprocess.CalledProcessError) as exc: + detail = getattr(exc, "stderr", None) or str(exc) + raise RuntimeError( + f"Git provenance command failed in {REPO_ROOT}: " + f"{' '.join(command)}: {str(detail).strip()}" + ) from exc return result.stdout.strip() diff --git a/src/sparsevllm/distributed/flashinfer_custom_all_reduce.py b/src/sparsevllm/distributed/flashinfer_custom_all_reduce.py new file mode 100644 index 00000000..9bc8761c --- /dev/null +++ b/src/sparsevllm/distributed/flashinfer_custom_all_reduce.py @@ -0,0 +1,192 @@ +from __future__ import annotations + +import ctypes +from contextlib import contextmanager +from typing import TYPE_CHECKING, Iterator + +import torch +import torch.distributed as dist + +from sparsevllm.utils.log import logger + +if TYPE_CHECKING: + from sparsevllm.distributed.parallel_context import ParallelGroup + + +class FlashInferCustomAllReduce: + """Small-tensor TP all-reduce using FlashInfer's vLLM kernel.""" + + def __init__( + self, + group: ParallelGroup, + *, + max_size_bytes: int = 8 * 1024 * 1024, + num_ctas: int = 32, + ) -> None: + if group.size != 2 or group.process_group is None: + raise ValueError( + "FlashInfer custom all-reduce requires a two-rank process group." + ) + if dist.get_backend(group.process_group) != dist.Backend.NCCL: + raise ValueError("FlashInfer custom all-reduce requires an NCCL group.") + if max_size_bytes <= 0 or num_ctas <= 0: + raise ValueError( + "FlashInfer custom all-reduce sizes must be positive, got " + f"max_size_bytes={max_size_bytes}, num_ctas={num_ctas}." + ) + if not all( + torch.cuda.can_device_access_peer(group.rank, peer_rank) + for peer_rank in range(group.size) + if peer_rank != group.rank + ): + raise RuntimeError( + "FlashInfer custom all-reduce requires CUDA peer access between " + f"all ranks in {group.ranks}." + ) + + from flashinfer.comm import ( + CudaRTLibrary, + create_shared_buffer, + vllm_all_reduce, + vllm_dispose, + vllm_get_graph_buffer_ipc_meta, + vllm_init_custom_ar, + vllm_meta_size, + vllm_register_buffer, + vllm_register_graph_buffers, + ) + + self.group = group + self.max_size_bytes = int(max_size_bytes) + self.num_ctas = int(num_ctas) + self._cudart = CudaRTLibrary() + self._cuda_ipc_close = self._cudart.lib.cudaIpcCloseMemHandle + self._cuda_ipc_close.restype = ctypes.c_int + self._cuda_ipc_close.argtypes = [ctypes.c_void_p] + self._all_reduce = vllm_all_reduce + self._dispose = vllm_dispose + self._get_graph_buffer_ipc_meta = vllm_get_graph_buffer_ipc_meta + self._register_graph_buffers = vllm_register_graph_buffers + self._capturing = False + self._closed = False + self._supported_sizes: dict[int, int] = {} + self._fallback_sizes: dict[int, int] = {} + self._captured_sizes: dict[int, int] = {} + + self._meta_ptrs = create_shared_buffer( + vllm_meta_size() + self.max_size_bytes, + group.process_group, + ) + self._buffer_ptrs = create_shared_buffer( + self.max_size_bytes, + group.process_group, + ) + self._rank_data = torch.empty( + 8 * 1024 * 1024, + dtype=torch.uint8, + device="cuda", + ) + # TP2 is valid without assuming NVLink. Peer access was checked above. + self._handle = vllm_init_custom_ar( + self._meta_ptrs, + self._rank_data, + group.rank, + False, + ) + vllm_register_buffer(self._handle, self._buffer_ptrs) + + def supports(self, tensor: torch.Tensor) -> bool: + size_bytes = int(tensor.numel() * tensor.element_size()) + return bool( + tensor.is_cuda + and tensor.dtype in {torch.bfloat16, torch.float16} + and tensor.is_contiguous() + and size_bytes > 0 + and size_bytes % 16 == 0 + and size_bytes <= self.max_size_bytes + ) + + def all_reduce(self, tensor: torch.Tensor) -> torch.Tensor | None: + size_bytes = int(tensor.numel() * tensor.element_size()) + if not self.supports(tensor): + self._fallback_sizes[size_bytes] = self._fallback_sizes.get(size_bytes, 0) + 1 + return None + self._supported_sizes[size_bytes] = self._supported_sizes.get(size_bytes, 0) + 1 + output = torch.empty_like(tensor) + registered = self._capturing and torch.cuda.is_current_stream_capturing() + if registered: + self._captured_sizes[size_bytes] = self._captured_sizes.get(size_bytes, 0) + 1 + self._all_reduce( + self._handle, + tensor, + output, + 0 if registered else self._buffer_ptrs[self.group.rank], + 0 if registered else self.max_size_bytes, + self.num_ctas, + ) + return output + + def _register_captured_buffers(self) -> None: + handles, offsets = self._get_graph_buffer_ipc_meta(self._handle) + metadata = [None] * self.group.size + dist.all_gather_object( + metadata, + [handles, offsets], + group=self.group.process_group, + ) + self._register_graph_buffers( + self._handle, + [item[0] for item in metadata], + [item[1] for item in metadata], + ) + logger.info( + "FlashInfer custom TP all-reduce capture evidence: rank={}, " + "captured_sizes={}, supported_sizes={}, fallback_sizes={}, " + "graph_buffers={}", + self.group.rank, + dict(sorted(self._captured_sizes.items())), + dict(sorted(self._supported_sizes.items())), + dict(sorted(self._fallback_sizes.items())), + len(offsets), + ) + + @contextmanager + def capture(self) -> Iterator[None]: + if self._capturing: + raise RuntimeError("FlashInfer custom all-reduce capture is already active.") + self._capturing = True + completed = False + try: + yield + completed = True + finally: + self._capturing = False + if completed: + self._register_captured_buffers() + + def close(self) -> None: + if self._closed: + return + self._closed = True + self._dispose(self._handle) + self._close_shared_buffer(self._buffer_ptrs) + self._close_shared_buffer(self._meta_ptrs) + + def _close_shared_buffer(self, pointers: list[int]) -> None: + device_ids = [torch.cuda.current_device()] + dist.barrier(group=self.group.process_group, device_ids=device_ids) + for peer_rank, pointer in enumerate(pointers): + if peer_rank == self.group.rank: + continue + result = int(self._cuda_ipc_close(ctypes.c_void_p(pointer))) + if result != 0: + raise RuntimeError( + "cudaIpcCloseMemHandle failed: " + f"{self._cudart.cudaGetErrorString(result)}" + ) + dist.barrier(group=self.group.process_group, device_ids=device_ids) + self._cudart.cudaFree(ctypes.c_void_p(pointers[self.group.rank])) + dist.barrier(group=self.group.process_group, device_ids=device_ids) + + +__all__ = ["FlashInferCustomAllReduce"] diff --git a/src/sparsevllm/distributed/parallel_context.py b/src/sparsevllm/distributed/parallel_context.py index eb41ae37..acacadca 100644 --- a/src/sparsevllm/distributed/parallel_context.py +++ b/src/sparsevllm/distributed/parallel_context.py @@ -1,6 +1,8 @@ from __future__ import annotations +from contextlib import contextmanager from dataclasses import dataclass, field, replace +from typing import Iterator, Protocol import torch import torch.distributed as dist @@ -9,6 +11,16 @@ from sparsevllm.operators.all_reduce import AllReduceProvider, resolve_all_reduce_provider +class GraphAllReduceProvider(Protocol): + group: ParallelGroup + + def all_reduce(self, tensor: torch.Tensor) -> torch.Tensor | None: ... + + def capture(self): ... + + def close(self) -> None: ... + + @dataclass(frozen=True) class ParallelGroup: process_group: dist.ProcessGroup | None @@ -35,6 +47,7 @@ class ParallelContext: expert: ParallelGroup data: ParallelGroup moe_tensor: ParallelGroup | None = None + all_reduce_provider: GraphAllReduceProvider | None = None @property def world_rank(self) -> int: @@ -122,6 +135,50 @@ def tp_all_reduce( ) -> torch.Tensor: return self.attention_tp_all_reduce(tensor, op) + def _all_reduce_out_of_place( + self, + tensor: torch.Tensor, + group: ParallelGroup, + op: dist.ReduceOp, + ) -> torch.Tensor: + provider = self.all_reduce_provider + if ( + provider is not None + and provider.group.ranks == group.ranks + and op == dist.ReduceOp.SUM + ): + output = provider.all_reduce(tensor) + if output is not None: + return output + return self._all_reduce(tensor, group, op) + + def tp_all_reduce_out_of_place( + self, + tensor: torch.Tensor, + op: dist.ReduceOp = dist.ReduceOp.SUM, + ) -> torch.Tensor: + return self._all_reduce_out_of_place(tensor, self.attention, op) + + def world_all_reduce_out_of_place( + self, + tensor: torch.Tensor, + op: dist.ReduceOp = dist.ReduceOp.SUM, + ) -> torch.Tensor: + return self._all_reduce_out_of_place(tensor, self.world, op) + + @contextmanager + def all_reduce_capture(self) -> Iterator[None]: + provider = self.all_reduce_provider + if provider is None: + yield + return + with provider.capture(): + yield + + def close_all_reduce_provider(self) -> None: + if self.all_reduce_provider is not None: + self.all_reduce_provider.close() + def ep_all_reduce( self, tensor: torch.Tensor, @@ -204,6 +261,7 @@ def _local_group( def init_parallel_context( *, topology: ParallelTopology, + enable_flashinfer_custom_all_reduce: bool = False, ) -> ParallelContext: global _PARALLEL_CONTEXT if _PARALLEL_CONTEXT is not None: @@ -267,6 +325,15 @@ def bind_provider(group: ParallelGroup) -> ParallelGroup: data=bind_provider(context.data), moe_tensor=bind_provider(context.moe_tensor or context.tensor), ) + if enable_flashinfer_custom_all_reduce: + from sparsevllm.distributed.flashinfer_custom_all_reduce import ( + FlashInferCustomAllReduce, + ) + + _PARALLEL_CONTEXT = replace( + _PARALLEL_CONTEXT, + all_reduce_provider=FlashInferCustomAllReduce(_PARALLEL_CONTEXT.tensor), + ) return _PARALLEL_CONTEXT diff --git a/src/sparsevllm/engine/decode_cuda_graph.py b/src/sparsevllm/engine/decode_cuda_graph.py index 1620262b..f0b6a963 100644 --- a/src/sparsevllm/engine/decode_cuda_graph.py +++ b/src/sparsevllm/engine/decode_cuda_graph.py @@ -88,6 +88,7 @@ def __init__( is_long_text_batch: Callable[[list[Sequence], bool], bool], method: str, rank: int, + parallel_context, capture_sizes: list[int], context_sizes: list[int] | tuple[int, ...] | str | int | None = None, graph_pool=None, @@ -100,6 +101,7 @@ def __init__( self.is_long_text_batch = is_long_text_batch self.method = str(method or "") self.rank = int(rank) + self.parallel_context = parallel_context self.platform = platforms.current_platform self.capture_sizes = sorted(set(int(size) for size in capture_sizes)) if not self.capture_sizes or any(size <= 0 for size in self.capture_sizes): @@ -438,15 +440,16 @@ def _capture( graph_input_sparse_state_refs = self._snapshot_sparse_state_refs() graph = torch.cuda.CUDAGraph() try: - with torch.cuda.graph(graph, pool=self.graph_pool): - self._reset_graph_input_attn_scores(graph_input_sparse_state_refs) - logits = self.run_model(input_ids, positions, is_prefill=False) - if state.key.capture_sampling: - if logits is None: - raise RuntimeError("decode_cuda_graph capture_sampling requires rank-0 logits.") - token_ids = logits.argmax(dim=-1) - else: - token_ids = None + with self.parallel_context.all_reduce_capture(): + with torch.cuda.graph(graph, pool=self.graph_pool): + self._reset_graph_input_attn_scores(graph_input_sparse_state_refs) + logits = self.run_model(input_ids, positions, is_prefill=False) + if state.key.capture_sampling: + if logits is None: + raise RuntimeError("decode_cuda_graph capture_sampling requires rank-0 logits.") + token_ids = logits.argmax(dim=-1) + else: + token_ids = None except Exception as exc: raise RuntimeError(f"decode_cuda_graph capture failed: {exc!r}") from exc diff --git a/src/sparsevllm/engine/model_runner.py b/src/sparsevllm/engine/model_runner.py index 036f3d4e..23a3f70e 100644 --- a/src/sparsevllm/engine/model_runner.py +++ b/src/sparsevllm/engine/model_runner.py @@ -149,8 +149,20 @@ def __init__( ) self.parallel_context = init_parallel_context( topology=config.parallel_topology, + enable_flashinfer_custom_all_reduce=( + hf_config.model_type == "glm4_moe_lite" + and bool(config.decode_cuda_graph) + and config.tensor_parallel_size == 2 + and config.expert_parallel_size == 1 + and config.data_parallel_size == 1 + and self.platform.is_cuda() + and self.platform.get_device_caps(rank).compute_capability == (9, 0) + ), ) - + if self.parallel_context.all_reduce_provider is not None: + logger.info( + "Enabled FlashInfer custom TP all-reduce for tensors up to 8 MiB." + ) # CUDA allocator peaks are process-global and survive LLMEngine.exit(). # Start a new lifecycle before model construction so KV sizing observes # only this engine's model load and persistent allocations. @@ -306,6 +318,7 @@ def __init__( is_long_text_batch=self._is_long_text_batch, method=self.config.vllm_sparse_method, rank=self.rank, + parallel_context=self.parallel_context, capture_sizes=decode_static_capture_sizes, context_sizes=decode_static_context_sizes, graph_pool=self.cuda_graph_pool, @@ -346,6 +359,7 @@ def exit(self): ) if self.rank == 0: self.shm.unlink() + self.parallel_context.close_all_reduce_provider() reset_parallel_context() dist.destroy_process_group() diff --git a/src/sparsevllm/layers/embed_head.py b/src/sparsevllm/layers/embed_head.py index 8df9a445..e8645acd 100644 --- a/src/sparsevllm/layers/embed_head.py +++ b/src/sparsevllm/layers/embed_head.py @@ -56,7 +56,7 @@ def forward(self, x: torch.Tensor): y = F.embedding(x, self.weight) if self.tp_size > 1: y = mask.unsqueeze(1) * y - return self.parallel_context.tp_all_reduce(y) + return self.parallel_context.tp_all_reduce_out_of_place(y) class ParallelLMHead(VocabParallelEmbedding): diff --git a/src/sparsevllm/layers/linear.py b/src/sparsevllm/layers/linear.py index 018c9a03..1384cc3e 100755 --- a/src/sparsevllm/layers/linear.py +++ b/src/sparsevllm/layers/linear.py @@ -188,6 +188,43 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return F.linear(x, self.weight, self.bias) +class MergedReplicatedLinear(ReplicatedLinear): + """One replicated GEMM loaded from multiple output-row shards.""" + + def __init__( + self, + input_size: int, + output_sizes: list[int], + bias: bool = False, + ): + if not output_sizes or any(int(size) <= 0 for size in output_sizes): + raise ValueError(f"output_sizes must be positive, got {output_sizes}.") + self.output_sizes = [int(size) for size in output_sizes] + super().__init__(input_size, sum(self.output_sizes), bias) + + def weight_loader( + self, + param: nn.Parameter, + loaded_weight: torch.Tensor, + loaded_shard_id: int, + ) -> None: + loaded_shard_id = int(loaded_shard_id) + if loaded_shard_id < 0 or loaded_shard_id >= len(self.output_sizes): + raise ValueError( + f"loaded_shard_id must be in [0, {len(self.output_sizes)}), " + f"got {loaded_shard_id}." + ) + shard_offset = sum(self.output_sizes[:loaded_shard_id]) + shard_size = self.output_sizes[loaded_shard_id] + target = param.data.narrow(0, shard_offset, shard_size) + if tuple(target.shape) != tuple(loaded_weight.shape): + raise ValueError( + "MergedReplicatedLinear shard shape mismatch: " + f"expected={tuple(target.shape)}, got={tuple(loaded_weight.shape)}." + ) + target.copy_(loaded_weight) + + class ColumnParallelLinear(LinearBase): def __init__( @@ -454,4 +491,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: y = self.quant_provider(x, self.weight, self.weight_scale_inv, bias) else: y = F.linear(x, self.weight, bias) - return self.parallel_context.tp_all_reduce(y) if self.reduce_results else y + if self.reduce_results: + return self.parallel_context.tp_all_reduce_out_of_place(y) + return y diff --git a/src/sparsevllm/layers/mla_attention.py b/src/sparsevllm/layers/mla_attention.py index 1d92826e..9da76fa5 100644 --- a/src/sparsevllm/layers/mla_attention.py +++ b/src/sparsevllm/layers/mla_attention.py @@ -98,6 +98,7 @@ class _MlaPrefillQueryPlan: validation_scope: object source_context_lens: torch.Tensor query_tokens: int + max_query_len: int def matches( self, @@ -329,6 +330,10 @@ def bind( def device(self) -> torch.device: return torch.device(getattr(self.provider, "device")) + @property + def supports_explicit_prefill(self) -> bool: + return bool(getattr(self.provider, "supports_explicit_prefill", False)) + def _require_mla_payload( self, view: PrefillComputeView | DecodeComputeView | AttentionKeyComputeView, @@ -721,6 +726,8 @@ def run_prefill( f"q={q.device}/{q.dtype} expected=" f"{self.device}/{self.spec.activation_dtype}." ) + query_tokens = int(q.shape[0]) + validation_scope = get_context().attention_validation_scope batch_size = int(history.context_lens.numel()) for name, tensor in ( ("b_start_loc", b_start_loc), @@ -736,8 +743,6 @@ def run_prefill( f"{name} must be int32 on {self.device}, got " f"{tensor.device}/{tensor.dtype}." ) - validation_scope = get_context().attention_validation_scope - query_tokens = int(q.shape[0]) cached_query_plan = self._prefill_query_plan if cached_query_plan is None or not cached_query_plan.matches( validation_scope, @@ -770,9 +775,33 @@ def run_prefill( validation_scope=validation_scope, source_context_lens=history.context_lens, query_tokens=query_tokens, + max_query_len=max(chunks), ) explicit_view = self.build_prefill_explicit_view(workset) + if self.supports_explicit_prefill: + run_prefill = getattr(self.provider, "run_explicit_prefill", None) + if not callable(run_prefill): + raise RuntimeError( + f"MLA provider {self.provider.name!r} advertises explicit " + "prefill without a run_explicit_prefill implementation." + ) + cu_seqlens_q = get_context().cu_seqlens_q + if cu_seqlens_q is None: + raise RuntimeError("MLA explicit prefill requires cu_seqlens_q.") + output = torch.empty( + (query_tokens, self.spec.local_q_heads, self.spec.value_head_dim), + dtype=q.dtype, + device=q.device, + ) + return run_prefill( + q, + explicit_view, + output, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=self._prefill_query_plan.max_query_len, + validation_scope=validation_scope, + ) return self.prefill_backend.run_prefill( q, explicit_view, @@ -827,7 +856,6 @@ def run_decode( validation_scope=get_context().attention_validation_scope, ) - __all__ = [ "MLAAttention", "MlaPrefillHistory", diff --git a/src/sparsevllm/models/glm4_moe_lite.py b/src/sparsevllm/models/glm4_moe_lite.py index 966de03b..873bf941 100644 --- a/src/sparsevllm/models/glm4_moe_lite.py +++ b/src/sparsevllm/models/glm4_moe_lite.py @@ -14,11 +14,15 @@ AttentionKeyComputeView, MlaLatentWrite, ) -from sparsevllm.layers.embed_head import ParallelLMHead, VocabParallelEmbedding +from sparsevllm.layers.embed_head import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from sparsevllm.layers.activation import SiluAndMul from sparsevllm.layers.layernorm import RMSNorm from sparsevllm.layers.linear import ( ColumnParallelLinear, - ReplicatedLinear, + MergedReplicatedLinear, RowParallelLinear, ) from sparsevllm.layers.mla_attention import MLAAttention @@ -26,7 +30,11 @@ from sparsevllm.layers.rotary_embedding import RotaryEmbedding, get_rope from sparsevllm.models.qwen3 import Qwen3MLP from sparsevllm.operators.mla_attention import MlaAttentionOpSpec -from sparsevllm.operators.moe import model_activation_dtype, resolve_moe_provider +from sparsevllm.operators.moe import ( + MoeOpSpec, + model_activation_dtype, + resolve_moe_provider, +) from sparsevllm.platforms import device_runtime from sparsevllm.utils.context import get_context from sparsevllm.utils.log import logger @@ -41,6 +49,10 @@ r"^model\.layers\.(\d+)\.mlp\.experts\.(\d+)\." r"(gate_proj|up_proj|down_proj)\.expert_weight$" ) +_SHARED_EXPERT_SOURCE_RE = re.compile( + r"^model\.layers\.(\d+)\.mlp\.shared_experts\." + r"(gate_proj|up_proj|down_proj)\.weight$" +) _MTP_LAYER_INDEX = 47 _MTP_PREFIX = f"model.layers.{_MTP_LAYER_INDEX}." @@ -120,9 +132,12 @@ def __init__( if bool(getattr(quantization, "enabled", False)): raise NotImplementedError("GLM MLA projections do not support quantization.") - self.q_a_proj = ReplicatedLinear( + self.fused_qkv_a_proj = MergedReplicatedLinear( int(config.hidden_size), - self.q_lora_rank, + [ + self.q_lora_rank, + self.kv_lora_rank + self.qk_rope_head_dim, + ], bias=bool(config.attention_bias), ) self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps) @@ -131,11 +146,6 @@ def __init__( self.num_heads * self.qk_head_dim, bias=False, ) - self.kv_a_proj_with_mqa = ReplicatedLinear( - int(config.hidden_size), - self.kv_lora_rank + self.qk_rope_head_dim, - bias=bool(config.attention_bias), - ) self.kv_a_layernorm = RMSNorm( self.kv_lora_rank, eps=config.rms_norm_eps, @@ -268,14 +278,22 @@ def _run_attention( [self.qk_nope_head_dim, self.v_head_dim], dim=-1, ) - expanded_rope = history.gathered_rope[:, None, :].expand( - -1, - self.local_heads, - -1, + expanded_k = torch.empty( + ( + history.visible_tokens, + self.local_heads, + self.qk_head_dim, + ), + dtype=expanded_k_nope.dtype, + device=expanded_k_nope.device, + ) + expanded_k[..., : self.qk_nope_head_dim].copy_(expanded_k_nope) + expanded_k[..., self.qk_nope_head_dim :].copy_( + history.gathered_rope[:, None, :] ) workset = self.mla_attention.bind_prefill_kv( history, - expanded_k=torch.cat((expanded_k_nope, expanded_rope), dim=-1), + expanded_k=expanded_k, expanded_v=expanded_v, ) b_start_loc = context.cu_seqlens_q[:-1] @@ -330,14 +348,18 @@ def forward( hidden_states: torch.Tensor, rotary_emb: RotaryEmbedding, ) -> torch.Tensor: - q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) + compressed_qkv = self.fused_qkv_a_proj(hidden_states) + compressed_q, compressed_kv = compressed_qkv.split( + [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], + dim=-1, + ) + q = self.q_b_proj(self.q_a_layernorm(compressed_q)) q = q.view(-1, self.local_heads, self.qk_head_dim) q_nope, q_rope = q.split( [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1, ) - compressed_kv = self.kv_a_proj_with_mqa(hidden_states) latent, k_rope = compressed_kv.split( [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1, @@ -377,23 +399,33 @@ def __init__(self, config: Glm4MoeLiteConfig) -> None: self.e_score_correction_bias = nn.Parameter( torch.empty(self.num_experts, dtype=torch.float32) ) - from sparsevllm.triton_kernel.moe_biased_sigmoid import ( - topk_biased_sigmoid, - ) + from sparsevllm.operators.sgl_moe import SglGlmFusedMoeGate - self.topk_impl = topk_biased_sigmoid + self.topk_impl = SglGlmFusedMoeGate( + num_experts=self.num_experts, + top_k=self.top_k, + ) + self._fused_topk_impl = self.topk_impl def forward( self, hidden_states: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: router_logits = F.linear(hidden_states.float(), self.weight) - topk_weights, topk_ids = self.topk_impl( - router_logits, - self.e_score_correction_bias, - top_k=self.top_k, - ) - topk_weights = topk_weights * self.routed_scaling_factor + if self.topk_impl is self._fused_topk_impl: + topk_weights, topk_ids = self.topk_impl( + router_logits, + self.e_score_correction_bias, + top_k=self.top_k, + routed_scaling_factor=self.routed_scaling_factor, + ) + else: + topk_weights, topk_ids = self.topk_impl( + router_logits, + self.e_score_correction_bias, + top_k=self.top_k, + ) + topk_weights = topk_weights * self.routed_scaling_factor return router_logits, topk_weights, topk_ids @@ -404,20 +436,114 @@ def __init__( *, decode_cuda_graph: bool, ) -> None: + parallel_context = get_parallel_context() + self.routed_num_experts = int(config.n_routed_experts) + self.routed_top_k = int(config.num_experts_per_tok) + self.fuses_shared_decode = bool( + decode_cuda_graph + and ( + self.routed_num_experts, + int(config.n_shared_experts), + self.routed_top_k, + int(config.hidden_size), + int(config.moe_intermediate_size), + int(parallel_context.moe_tp_size), + int(parallel_context.ep_size), + ) + == (64, 1, 4, 2048, 1536, 2, 1) + ) + packed_num_experts = self.routed_num_experts + int( + self.fuses_shared_decode + ) + packed_top_k = self.routed_top_k + int(self.fuses_shared_decode) super().__init__( - num_experts=int(config.n_routed_experts), + num_experts=packed_num_experts, hidden_size=int(config.hidden_size), intermediate_size=int(config.moe_intermediate_size), - top_k=int(config.num_experts_per_tok), + top_k=packed_top_k, activation_dtype=model_activation_dtype(config), fp8_enabled=False, cuda_graph=bool(decode_cuda_graph), routing_method="biased_sigmoid", model_label="GLM-4.7-Flash", provider_resolver=resolve_moe_provider, - parallel_context=get_parallel_context(), + parallel_context=parallel_context, + ) + self.shared_expert_id = ( + self.routed_num_experts if self.fuses_shared_decode else None + ) + self.shared_act = SiluAndMul() + if self.fuses_shared_decode: + self.routed_op_spec = MoeOpSpec( + num_experts=self.routed_num_experts, + num_local_experts=self.routed_num_experts, + hidden_size=self.hidden_size, + intermediate_size=self.intermediate_size, + top_k=self.routed_top_k, + activation_dtype=self.op_spec.activation_dtype, + weight_dtype=self.op_spec.weight_dtype, + block_shape=self.op_spec.block_shape, + ep_size=1, + cuda_graph=self.op_spec.cuda_graph, + tp_size=self.op_spec.tp_size, + routing_method=self.op_spec.routing_method, + scale_dtype=self.op_spec.scale_dtype, + ) + self.routed_provider = resolve_moe_provider(self.routed_op_spec) + else: + self.routed_op_spec = self.op_spec + self.routed_provider = self.provider + + def forward( + self, + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + ) -> torch.Tensor: + if not self.fuses_shared_decode: + return super().forward(hidden_states, topk_ids, topk_weights) + return self.routed_provider.run( + self.routed_op_spec, + hidden_states, + topk_ids, + topk_weights, + self.w13_weight[: self.routed_num_experts], + self.w2_weight[: self.routed_num_experts], + self.w13_scale_inv, + self.w2_scale_inv, + local_expert_start=0, + ep_rank=int(self.ep_rank), + ) + + def forward_shared(self, hidden_states: torch.Tensor) -> torch.Tensor: + if self.shared_expert_id is None: + raise RuntimeError("Packed shared expert is not enabled.") + gate_up = F.linear( + hidden_states, + self.w13_weight[self.shared_expert_id], + ) + return F.linear( + self.shared_act(gate_up), + self.w2_weight[self.shared_expert_id], ) + def forward_routed_and_shared( + self, + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + ) -> torch.Tensor: + if self.shared_expert_id is None: + raise RuntimeError("Packed shared expert is not enabled.") + from sparsevllm.triton_kernel.moe import append_shared_expert_route + + fused_ids, fused_weights = append_shared_expert_route( + topk_ids, + topk_weights, + shared_expert_id=self.shared_expert_id, + ) + return super().forward(hidden_states, fused_ids, fused_weights) + class Glm4MoeLiteSparseMoeBlock(nn.Module): def __init__( @@ -439,20 +565,42 @@ def __init__( config, decode_cuda_graph=decode_cuda_graph, ) - self.shared_experts = Qwen3MLP( - hidden_size=int(config.hidden_size), - intermediate_size=( - int(config.moe_intermediate_size) * int(config.n_shared_experts) - ), - hidden_act=str(config.hidden_act), - mlp_chunk_size=self.mlp_chunk_size, - quantization=None, + self.shared_experts = ( + None + if self.experts.fuses_shared_decode + else Qwen3MLP( + hidden_size=int(config.hidden_size), + intermediate_size=( + int(config.moe_intermediate_size) + * int(config.n_shared_experts) + ), + hidden_act=str(config.hidden_act), + mlp_chunk_size=self.mlp_chunk_size, + quantization=None, + reduce_results=self.parallel_context.ep_size > 1, + ) ) def _routed_chunk(self, hidden_states: torch.Tensor) -> torch.Tensor: _, topk_weights, topk_ids = self.gate(hidden_states) return self.experts(hidden_states, topk_ids, topk_weights) + def _shared_chunk(self, hidden_states: torch.Tensor) -> torch.Tensor: + if self.shared_experts is not None: + return self.shared_experts(hidden_states) + return self.experts.forward_shared(hidden_states) + + def _routed_and_shared_chunk( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + _, topk_weights, topk_ids = self.gate(hidden_states) + return self.experts.forward_routed_and_shared( + hidden_states, + topk_ids, + topk_weights, + ) + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: if hidden_states.ndim != 2: raise ValueError( @@ -462,6 +610,32 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: debug_enabled = os.getenv("SPARSEVLLM_DEBUG_MOE", "0") == "1" if debug_enabled: self.debug_last_input = hidden_states.detach().clone() + context = get_context() + if ( + not debug_enabled + and getattr( + getattr(self, "experts", None), + "fuses_shared_decode", + False, + ) + and not context.is_prefill + ): + if int(hidden_states.shape[0]) <= self.mlp_chunk_size: + local_output = self._routed_and_shared_chunk(hidden_states) + else: + local_output = torch.cat( + [ + self._routed_and_shared_chunk(chunk) + for chunk in hidden_states.split( + self.mlp_chunk_size, + dim=0, + ) + ], + dim=0, + ) + return self.parallel_context.world_all_reduce_out_of_place( + local_output + ) if not debug_enabled: if int(hidden_states.shape[0]) <= self.mlp_chunk_size: routed = self._routed_chunk(hidden_states) @@ -507,21 +681,42 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: local_hit_count = local_mask.sum() self.debug_last_local_hit_count = ( local_hit_count - if torch.cuda.is_available() - and torch.cuda.is_current_stream_capturing() + if device_runtime.is_stream_capturing() else int(local_hit_count.item()) ) - # Each rank owns one (expert, intermediate) shard. Summing across the - # world composes pure TP, pure EP, and outer-TP x MoE-EP identically to - # the shared PackedMoeExperts contract used by Qwen3-MoE. - self.parallel_context.world_all_reduce(routed) if debug_enabled: + # Preserve the general EP composition and routed-only debug + # evidence. The single-reduction specialization below is TP-only. + self.parallel_context.world_all_reduce(routed) self.debug_last_routed_output = routed.detach().clone() - output = routed + self.shared_experts(hidden_states) + shared = self._shared_chunk(hidden_states) + if self.parallel_context.ep_size == 1: + shared = self.parallel_context.tp_all_reduce(shared) + output = routed + shared + elif self.parallel_context.ep_size > 1: + self.parallel_context.world_all_reduce(routed) + output = routed + self._shared_chunk(hidden_states) + else: + shared_local = self._shared_chunk(hidden_states) + if context.is_prefill: + # Both branches are TP partials. Compose them locally so the + # full MoE block needs one collective, matching the fused-MoE + # communication contract used by the reference runtime. + output = self.parallel_context.world_all_reduce_out_of_place( + routed + shared_local + ) + else: + # Decode tensors are small. Pack both partials into one + # collective, then add the independently reduced rows. This + # preserves the original BF16 reduction/addition order. + partials = self.parallel_context.world_all_reduce_out_of_place( + torch.stack((routed, shared_local), dim=0) + ) + output = partials[0] + partials[1] if debug_enabled: # ModelRunner's cross-rank evidence contract consumes the final # MoE block output, including both the synced routed experts and - # the replicated shared experts. + # the synchronized shared-expert branch. self.debug_last_output = output.detach().clone() return output @@ -728,6 +923,8 @@ def _expected_mtp_weight_names(num_experts: int) -> set[str]: class Glm4MoeLiteForCausalLM(nn.Module): special_weight_loaders = (".expert_weight",) packed_modules_mapping = { + "self_attn.q_a_proj": ("self_attn.fused_qkv_a_proj", 0), + "self_attn.kv_a_proj_with_mqa": ("self_attn.fused_qkv_a_proj", 1), "gate_proj": ("gate_up_proj", 0), "up_proj": ("gate_up_proj", 1), } @@ -830,7 +1027,7 @@ def warmup_moe(self, num_tokens: int = 1) -> None: dtype=torch.int64, device=hidden_states.device, ) - .remainder(experts.num_local_experts) + .remainder(experts.routed_num_experts) .add(experts.local_expert_start) .view(num_tokens, top_k) ) @@ -846,6 +1043,15 @@ def warmup_moe(self, num_tokens: int = 1) -> None: def map_weight_name(self, source_weight_name: str) -> str | None: if source_weight_name.startswith(_MTP_PREFIX): return None + shared_match = _SHARED_EXPERT_SOURCE_RE.match(source_weight_name) + if shared_match is not None: + layer_idx, projection = shared_match.groups() + experts = self._sparse_block(int(layer_idx)).experts + if experts.shared_expert_id is not None: + return ( + f"model.layers.{layer_idx}.mlp.experts." + f"{experts.shared_expert_id}.{projection}.expert_weight" + ) match = _EXPERT_SOURCE_RE.match(source_weight_name) if match is None: return source_weight_name diff --git a/src/sparsevllm/models/qwen3.py b/src/sparsevllm/models/qwen3.py index 9f467fe1..898d3116 100755 --- a/src/sparsevllm/models/qwen3.py +++ b/src/sparsevllm/models/qwen3.py @@ -169,6 +169,7 @@ def __init__( hidden_act: str, mlp_chunk_size: int = 16384, quantization=None, + reduce_results: bool = True, ) -> None: super().__init__() self.gate_up_proj = MergedColumnParallelLinear( @@ -182,6 +183,7 @@ def __init__( hidden_size, bias=False, quantization=quantization, + reduce_results=reduce_results, ) assert hidden_act == "silu" self.act_fn = SiluAndMul() diff --git a/src/sparsevllm/operators/mla_attention.py b/src/sparsevllm/operators/mla_attention.py index d5c6aa86..d6f79987 100644 --- a/src/sparsevllm/operators/mla_attention.py +++ b/src/sparsevllm/operators/mla_attention.py @@ -7,19 +7,24 @@ import sparsevllm.platforms as platforms from sparsevllm.engine.cache_manager.base import ( DecodeComputeView, + ExplicitKVPayload, MlaLatentPayload, + PrefillComputeView, ) from sparsevllm.operators.registry import ( OpRegistry, OpResolver, SupportResult, ) +from sparsevllm.operators.sgl_fa3 import SglFa3DecodeKernel, sgl_fa3_support from sparsevllm.platforms.interface import DeviceCaps, PlatformEnum from sparsevllm.triton_kernel.mla import ( DEFAULT_GLM_MLA_DECODE_CONFIG, + GLM_MLA_MAX_WORKSPACE_CONFIG, MlaDecodeLaunchConfig, allocate_mla_decode_workspace, run_mla_decode, + select_glm_mla_decode_config, validate_mla_decode_metadata, ) @@ -76,6 +81,7 @@ def softmax_scale(self) -> float: class MlaAttentionProvider: name = "" priority = 0 + supports_explicit_prefill = False def run( self, @@ -108,7 +114,7 @@ def __init__( op_spec: MlaAttentionOpSpec, device: torch.device | str, max_batch_size: int, - launch_config: MlaDecodeLaunchConfig = DEFAULT_GLM_MLA_DECODE_CONFIG, + launch_config: MlaDecodeLaunchConfig | None = None, ) -> None: self.spec = op_spec requested_device = torch.device(device) @@ -118,12 +124,14 @@ def __init__( "MLA max_batch_size must be positive, got " f"{self.max_batch_size}." ) - self.launch_config = launch_config + self._fixed_launch_config = launch_config + self.launch_config = launch_config or DEFAULT_GLM_MLA_DECODE_CONFIG + workspace_config = launch_config or GLM_MLA_MAX_WORKSPACE_CONFIG self.workspace = allocate_mla_decode_workspace( batch_size=self.max_batch_size, head_count=self.spec.local_q_heads, device=requested_device, - config=self.launch_config, + config=workspace_config, ) self.device = self.workspace.block_size.device self._validated_decode_metadata: tuple[ @@ -271,22 +279,13 @@ def _validate_run_inputs( ) return view.payload - @torch.no_grad() - def run( + def _validate_metadata( self, - q_nope_absorbed: torch.Tensor, - q_rope: torch.Tensor, - view: DecodeComputeView, - output: torch.Tensor, + view: DecodeComputeView | PrefillComputeView, + payload: MlaLatentPayload, *, - validation_scope: object | None = None, - ) -> torch.Tensor: - payload = self._validate_run_inputs( - q_nope_absorbed, - q_rope, - view, - output, - ) + validation_scope: object | None, + ) -> None: cache_slot_count = int(payload.latent_cache.shape[0]) metadata_key = ( view.meta.active_slots, @@ -311,22 +310,67 @@ def run( and entry[4] == metadata_key[4] for entry in cached_entries ) - if not metadata_is_validated: - validate_mla_decode_metadata( - view.meta.active_slots, - view.meta.req_indices, - view.meta.context_lens, - cache_slot_count=cache_slot_count, - max_context_len=view.meta.max_context_len, - ) - if validation_scope is None: - self._validated_decode_metadata = None - else: - cached_entries.append(metadata_key) - self._validated_decode_metadata = ( - validation_scope, - cached_entries, - ) + if metadata_is_validated: + return + validate_mla_decode_metadata( + view.meta.active_slots, + view.meta.req_indices, + view.meta.context_lens, + cache_slot_count=cache_slot_count, + max_context_len=view.meta.max_context_len, + ) + if validation_scope is None: + self._validated_decode_metadata = None + else: + cached_entries.append(metadata_key) + self._validated_decode_metadata = (validation_scope, cached_entries) + + def _launch_config_for( + self, + *, + batch_size: int, + max_context_len: int | None, + active_slot_width: int, + ) -> MlaDecodeLaunchConfig: + if self._fixed_launch_config is not None: + return self._fixed_launch_config + context_capacity = ( + active_slot_width + if max_context_len is None + else int(max_context_len) + ) + return select_glm_mla_decode_config( + batch_size=batch_size, + max_context_len=context_capacity, + local_q_heads=self.spec.local_q_heads, + ) + + @torch.no_grad() + def run( + self, + q_nope_absorbed: torch.Tensor, + q_rope: torch.Tensor, + view: DecodeComputeView, + output: torch.Tensor, + *, + validation_scope: object | None = None, + ) -> torch.Tensor: + payload = self._validate_run_inputs( + q_nope_absorbed, + q_rope, + view, + output, + ) + self._validate_metadata( + view, + payload, + validation_scope=validation_scope, + ) + launch_config = self._launch_config_for( + batch_size=int(q_nope_absorbed.shape[0]), + max_context_len=view.meta.max_context_len, + active_slot_width=int(view.meta.active_slots.shape[1]), + ) return run_mla_decode( q_nope_absorbed, q_rope, @@ -340,17 +384,194 @@ def run( softmax_scale=self.spec.softmax_scale, attn_score=view.meta.attn_score, max_context_len=view.meta.max_context_len, - config=self.launch_config, + config=launch_config, validate_metadata=False, ) +@MLA_ATTENTION_REGISTRY.register +class MlaSglFa3Provider(MlaTritonProvider): + """SGL FA3 decode with the score-producing Triton path kept explicit.""" + + name = "sgl_fa3_h100" + priority = 200 + supports_explicit_prefill = True + + def __init__( + self, + *, + op_spec: MlaAttentionOpSpec, + device: torch.device | str, + max_batch_size: int, + launch_config: MlaDecodeLaunchConfig | None = None, + ) -> None: + super().__init__( + op_spec=op_spec, + device=device, + max_batch_size=max_batch_size, + launch_config=launch_config, + ) + self.fa3 = SglFa3DecodeKernel( + device=self.device, + max_batch_size=self.max_batch_size, + softmax_scale=self.spec.softmax_scale, + ) + + @classmethod + def supports( + cls, + spec: MlaAttentionOpSpec, + caps: DeviceCaps, + ) -> SupportResult: + base = MlaTritonProvider.supports(spec, caps) + if not base.supported: + return base + supported, reason = sgl_fa3_support() + return SupportResult.yes(reason) if supported else SupportResult.no(reason) + + @torch.no_grad() + def run( + self, + q_nope_absorbed: torch.Tensor, + q_rope: torch.Tensor, + view: DecodeComputeView, + output: torch.Tensor, + *, + validation_scope: object | None = None, + ) -> torch.Tensor: + if view.meta.attn_score is not None: + return super().run( + q_nope_absorbed, + q_rope, + view, + output, + validation_scope=validation_scope, + ) + payload = self._validate_run_inputs( + q_nope_absorbed, + q_rope, + view, + output, + ) + self._validate_metadata( + view, + payload, + validation_scope=validation_scope, + ) + return self.fa3( + q_rope, + q_nope_absorbed, + payload.rope_cache, + payload.latent_cache, + view.meta.active_slots, + view.meta.req_indices, + view.meta.context_lens, + output, + num_splits=5, + validation_scope=validation_scope, + ) + + @torch.no_grad() + def run_explicit_prefill( + self, + q: torch.Tensor, + view: PrefillComputeView, + output: torch.Tensor, + *, + cu_seqlens_q: torch.Tensor, + max_seqlen_q: int, + validation_scope: object | None = None, + ) -> torch.Tensor: + if not isinstance(view, PrefillComputeView): + raise TypeError( + "MlaSglFa3Provider.run_explicit_prefill requires " + "PrefillComputeView, got " + f"{type(view).__name__}." + ) + if not isinstance(view.payload, ExplicitKVPayload): + raise TypeError( + "MLA explicit prefill requires ExplicitKVPayload, got " + f"{type(view.payload).__name__}." + ) + query_tokens = int(q.shape[0]) + expected_q_shape = ( + query_tokens, + self.spec.local_q_heads, + self.spec.qk_head_dim, + ) + if tuple(q.shape) != expected_q_shape: + raise ValueError( + f"q must have shape {expected_q_shape}, got {tuple(q.shape)}." + ) + expected_output_shape = ( + query_tokens, + self.spec.local_q_heads, + self.spec.value_head_dim, + ) + if tuple(output.shape) != expected_output_shape: + raise ValueError( + f"output must have shape {expected_output_shape}, got " + f"{tuple(output.shape)}." + ) + batch_size = int(view.meta.context_lens.numel()) + if batch_size > self.max_batch_size: + raise ValueError( + "MLA prefill batch exceeds provider capacity: " + f"batch={batch_size} max_batch_size={self.max_batch_size}." + ) + if cu_seqlens_q.shape != (batch_size + 1,): + raise ValueError( + f"cu_seqlens_q must have shape ({batch_size + 1},), got " + f"{tuple(cu_seqlens_q.shape)}." + ) + if cu_seqlens_q.device != self.device or cu_seqlens_q.dtype != torch.int32: + raise TypeError( + "cu_seqlens_q must be int32 on the provider device, got " + f"{cu_seqlens_q.device}/{cu_seqlens_q.dtype}." + ) + if not 0 < int(max_seqlen_q) <= query_tokens: + raise ValueError( + "max_seqlen_q must be in [1, query_tokens], got " + f"{max_seqlen_q} for {query_tokens}." + ) + payload = view.payload + tensors = { + "q": q, + "output": output, + "k_cache": payload.k_cache, + "v_cache": payload.v_cache, + } + for name, tensor in tensors.items(): + if tensor.device != self.device: + raise ValueError( + f"{name} is on {tensor.device}, expected {self.device}." + ) + expected_dtype = ( + self.spec.activation_dtype + ) + if tensor.dtype != expected_dtype: + raise TypeError( + f"{name} must use {expected_dtype}, got {tensor.dtype}." + ) + return self.fa3.run_explicit_varlen( + q, + payload.k_cache, + payload.v_cache, + view.meta.active_slots, + view.meta.req_indices, + view.meta.context_lens, + output, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=int(max_seqlen_q), + validation_scope=validation_scope, + ) + def resolve_mla_attention_provider( spec: MlaAttentionOpSpec, *, device: torch.device | str, max_batch_size: int, - launch_config: MlaDecodeLaunchConfig = DEFAULT_GLM_MLA_DECODE_CONFIG, + launch_config: MlaDecodeLaunchConfig | None = None, ) -> MlaAttentionProvider: """Resolve and bind an MLA provider during model construction.""" @@ -371,6 +592,7 @@ def resolve_mla_attention_provider( "MLA_ATTENTION_REGISTRY", "MlaAttentionOpSpec", "MlaAttentionProvider", + "MlaSglFa3Provider", "MlaTritonProvider", "resolve_mla_attention_provider", ] diff --git a/src/sparsevllm/operators/moe.py b/src/sparsevllm/operators/moe.py index 5ab1ae7e..5267675a 100644 --- a/src/sparsevllm/operators/moe.py +++ b/src/sparsevllm/operators/moe.py @@ -176,6 +176,77 @@ def run( MOE_REGISTRY: OpRegistry[MoeOpSpec, MoeProvider] = OpRegistry("routed MoE") +@MOE_REGISTRY.register +class SglAlignedTritonGlmMoeProvider(MoeProvider): + name = "sgl_aligned_triton_glm" + priority = 30 + gate_up_order = "gate_up" + + @classmethod + def supports(cls, spec: MoeOpSpec, caps: DeviceCaps) -> SupportResult: + if caps.platform != PlatformEnum.CUDA or caps.compute_capability != (9, 0): + return SupportResult.no("requires CUDA SM90") + if caps.device_name != "NVIDIA H100 80GB HBM3": + return SupportResult.no("requires profiled NVIDIA H100 80GB HBM3") + expected = { + (64, 64, 2048, 768, 4, 2, 1, "biased_sigmoid"), + (65, 65, 2048, 768, 5, 2, 1, "biased_sigmoid"), + } + actual = ( + spec.num_experts, + spec.num_local_experts, + spec.hidden_size, + spec.intermediate_size, + spec.top_k, + spec.tp_size, + spec.ep_size, + spec.routing_method, + ) + if actual not in expected: + return SupportResult.no( + f"requires a profiled GLM TP2 MoE shape {expected}, got {actual}" + ) + if spec.activation_dtype != torch.bfloat16: + return SupportResult.no("requires BF16 activations") + if spec.weight_dtype != torch.bfloat16 or spec.block_shape is not None: + return SupportResult.no("requires unquantized BF16 expert weights") + from sparsevllm.operators.sgl_moe import sgl_moe_alignment_support + + supported, reason = sgl_moe_alignment_support() + return SupportResult.yes() if supported else SupportResult.no(reason) + + def run( + self, + spec, + hidden_states, + topk_ids, + topk_weights, + w13_weight, + w2_weight, + w13_scale_inv, + w2_scale_inv, + *, + local_expert_start, + ep_rank, + ): + del ep_rank + if w13_scale_inv is not None or w2_scale_inv is not None: + raise RuntimeError("SGL-aligned BF16 MoE does not accept expert scales.") + from sparsevllm.operators.sgl_moe import sgl_moe_align_block_size + from sparsevllm.triton_kernel.moe import fused_moe + + return fused_moe( + hidden_states, + w13_weight, + w2_weight, + topk_ids, + topk_weights, + num_experts=spec.num_experts, + local_expert_start=local_expert_start, + alignment_impl=sgl_moe_align_block_size, + ) + + @MOE_REGISTRY.register class TritonMinimaxM2FusedMoeProvider(MoeProvider): name = "triton_minimax_m2_fused" diff --git a/src/sparsevllm/operators/sgl_fa3.py b/src/sparsevllm/operators/sgl_fa3.py new file mode 100644 index 00000000..f029cdfd --- /dev/null +++ b/src/sparsevllm/operators/sgl_fa3.py @@ -0,0 +1,510 @@ +from __future__ import annotations + +import importlib.metadata +import importlib.util +import re +from collections.abc import Sequence +from dataclasses import dataclass + +import torch + + +_SGL_KERNEL_DISTRIBUTIONS = ("sglang-kernel", "sgl-kernel") +_MIN_VALIDATED_VERSIONS = { + (0, 3): (0, 3, 14), + (0, 4): (0, 4, 5), +} + +_COMMON_FWD_ARGUMENTS = ( + "q", + "k", + "v", + "k_new", + "v_new", + "q_v", + "out", + "cu_seqlens_q", + "cu_seqlens_k", + "cu_seqlens_k_new", + "seqused_q", + "seqused_k", + "max_seqlen_q", + "max_seqlen_k", + "page_table", + "kv_batch_idx", + "leftpad_k", + "rotary_cos", + "rotary_sin", + "seqlens_rotary", + "q_descale", + "k_descale", + "v_descale", + "softmax_scale", + "is_causal", + "window_size_left", + "window_size_right", +) +_FWD_TRAILING_ARGUMENTS = ( + "softcap", + "is_rotary_interleaved", + "scheduler_metadata", + "num_splits", + "pack_gqa", + "sm_margin", + "sinks", +) +_FWD_04_EXTRA_ARGUMENTS = ( + "sparse_mask_fine", + "only_qv", +) + + +@dataclass(frozen=True, slots=True) +class _SchedulerMetadataPlan: + """Step-local FA3 scheduling metadata shared by every MLA layer.""" + + validation_scope: object + context_lens: torch.Tensor + cu_seqlens_q: torch.Tensor + batch_size: int + max_seqlen_q: int + max_seqlen_k: int + num_heads: int + num_heads_k: int + headdim: int + headdim_v: int + qkv_dtype: torch.dtype + num_splits: int + metadata: torch.Tensor + + def matches( + self, + *, + validation_scope: object, + context_lens: torch.Tensor, + cu_seqlens_q: torch.Tensor, + batch_size: int, + max_seqlen_q: int, + max_seqlen_k: int, + num_heads: int, + num_heads_k: int, + headdim: int, + headdim_v: int, + qkv_dtype: torch.dtype, + num_splits: int, + ) -> bool: + return ( + self.validation_scope is validation_scope + and self.context_lens.data_ptr() == context_lens.data_ptr() + and self.cu_seqlens_q.data_ptr() == cu_seqlens_q.data_ptr() + and self.batch_size == int(batch_size) + and self.max_seqlen_q == int(max_seqlen_q) + and self.max_seqlen_k == int(max_seqlen_k) + and self.num_heads == int(num_heads) + and self.num_heads_k == int(num_heads_k) + and self.headdim == int(headdim) + and self.headdim_v == int(headdim_v) + and self.qkv_dtype == qkv_dtype + and self.num_splits == int(num_splits) + ) + + +def _parse_version(version: str) -> tuple[int, int, int] | None: + parts = str(version).split(".") + if len(parts) < 3: + return None + parsed: list[int] = [] + for part in parts[:3]: + match = re.match(r"(\d+)", part) + if match is None: + return None + parsed.append(int(match.group(1))) + return tuple(parsed) # type: ignore[return-value] + + +def _installed_sgl_kernel_version() -> tuple[str, tuple[int, int, int]] | None: + for distribution in _SGL_KERNEL_DISTRIBUTIONS: + try: + version = importlib.metadata.version(distribution) + except importlib.metadata.PackageNotFoundError: + continue + parsed = _parse_version(version) + if parsed is None: + return version, (-1, -1, -1) + return version, parsed + return None + + +def sgl_fa3_support() -> tuple[bool, str]: + """Probe package metadata without importing optional SGL Python wrappers.""" + + try: + module_spec = importlib.util.find_spec("sgl_kernel") + except (ImportError, ValueError): + module_spec = None + if module_spec is None: + return False, "sgl-kernel is not installed" + installed = _installed_sgl_kernel_version() + if installed is None: + return False, "sgl-kernel package metadata is unavailable" + version, parsed = installed + minimum = _MIN_VALIDATED_VERSIONS.get(parsed[:2]) + if minimum is None: + return ( + False, + "requires validated sgl-kernel 0.3.x or sglang-kernel 0.4.x " + f"API, got {version}", + ) + if parsed < minimum: + package = "sglang-kernel" if parsed[:2] == (0, 4) else "sgl-kernel" + required = ".".join(str(value) for value in minimum) + return False, f"requires {package} >= {required}, got {version}" + return True, f"SGL kernel {version} FA3 raw op is available" + + +class SglFa3DecodeKernel: + """Allocation-free GLM MLA decode adapter for the SGL FA3 raw op.""" + + def __init__( + self, + *, + device: torch.device, + max_batch_size: int, + softmax_scale: float, + num_splits: int = 0, + ) -> None: + supported, reason = sgl_fa3_support() + if not supported: + raise RuntimeError(reason) + + # Importing flash_ops registers FA3 only. Importing + # sgl_kernel.flash_attn would also import the optional FA4 wrapper, + # whose CUTLASS Python dependency is unrelated to this provider. + from sgl_kernel import flash_ops as _flash_ops # noqa: F401 + + op = torch.ops.sgl_kernel.fwd.default + argument_names = tuple(argument.name for argument in op._schema.arguments) + legacy = _COMMON_FWD_ARGUMENTS + _FWD_TRAILING_ARGUMENTS + chunked = ( + _COMMON_FWD_ARGUMENTS + + ("attention_chunk",) + + _FWD_TRAILING_ARGUMENTS + ) + chunked_04 = chunked + _FWD_04_EXTRA_ARGUMENTS + if argument_names == legacy: + self._has_attention_chunk = False + self._has_04_arguments = False + elif argument_names == chunked: + self._has_attention_chunk = True + self._has_04_arguments = False + elif argument_names == chunked_04: + self._has_attention_chunk = True + self._has_04_arguments = True + else: + raise RuntimeError( + "Unsupported sgl-kernel FA3 fwd schema: " + f"arguments={argument_names}." + ) + self._op = op + scheduler_packet = getattr( + torch.ops.sgl_kernel, + "get_scheduler_metadata", + None, + ) + # sgl-kernel 0.3.14 exposes the FA3 scheduler only internally. Keep + # its validated per-call scheduling path; 0.4.5 exposes the reusable + # raw op used below. + self._scheduler_op = ( + None if scheduler_packet is None else scheduler_packet.default + ) + self._scheduler_plan: _SchedulerMetadataPlan | None = None + self._captured_scheduler_plans: list[_SchedulerMetadataPlan] = [] + self._cu_seqlens_q = torch.arange( + int(max_batch_size) + 1, + dtype=torch.int32, + device=device, + ) + self.softmax_scale = float(softmax_scale) + self.num_splits = int(num_splits) + if self.num_splits < 0: + raise ValueError( + f"FA3 num_splits must be non-negative, got {self.num_splits}." + ) + + def __call__( + self, + q_rope: torch.Tensor, + q_latent: torch.Tensor, + rope_cache: torch.Tensor, + latent_cache: torch.Tensor, + page_table: torch.Tensor, + request_indices: torch.Tensor, + context_lens: torch.Tensor, + output: torch.Tensor, + *, + num_splits: int | None = None, + validation_scope: object | None = None, + ) -> torch.Tensor: + batch_size = int(q_rope.shape[0]) + return self.run_varlen( + q_rope, + q_latent, + rope_cache, + latent_cache, + page_table, + request_indices, + context_lens, + output, + cu_seqlens_q=self._cu_seqlens_q[: batch_size + 1], + max_seqlen_q=1, + num_splits=num_splits, + validation_scope=validation_scope, + ) + + def _scheduler_metadata( + self, + q: torch.Tensor, + k_cache: torch.Tensor, + page_table: torch.Tensor, + context_lens: torch.Tensor, + cu_seqlens_q: torch.Tensor, + *, + headdim_v: int, + max_seqlen_q: int, + num_splits: int, + validation_scope: object | None, + ) -> torch.Tensor | None: + if self._scheduler_op is None: + return None + batch_size = int(context_lens.numel()) + max_seqlen_k = int(page_table.shape[1]) + num_heads = int(q.shape[1]) + num_heads_k = int(k_cache.shape[1]) + headdim = int(q.shape[-1]) + headdim_v = int(headdim_v) + is_capturing = torch.cuda.is_current_stream_capturing() + plans = ( + reversed(self._captured_scheduler_plans) + if is_capturing + else (self._scheduler_plan,) + ) + plan = next( + ( + candidate + for candidate in plans + if validation_scope is not None + and candidate is not None + and candidate.matches( + validation_scope=validation_scope, + context_lens=context_lens, + cu_seqlens_q=cu_seqlens_q, + batch_size=batch_size, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + num_heads=num_heads, + num_heads_k=num_heads_k, + headdim=headdim, + headdim_v=headdim_v, + qkv_dtype=q.dtype, + num_splits=num_splits, + ) + ), + None, + ) + if plan is not None: + return plan.metadata + + metadata = self._scheduler_op( + batch_size, + int(max_seqlen_q), + max_seqlen_k, + num_heads, + num_heads_k, + headdim, + headdim_v, + q.dtype, + context_lens, + cu_seqlens_q, + None, + None, + None, + None, + 1, + 0, + True, + -1, + -1, + 0, + False, + int(num_splits), + None, + 0, + ) + if validation_scope is not None: + plan = _SchedulerMetadataPlan( + validation_scope=validation_scope, + context_lens=context_lens, + cu_seqlens_q=cu_seqlens_q, + batch_size=batch_size, + max_seqlen_q=int(max_seqlen_q), + max_seqlen_k=max_seqlen_k, + num_heads=num_heads, + num_heads_k=num_heads_k, + headdim=headdim, + headdim_v=headdim_v, + qkv_dtype=q.dtype, + num_splits=int(num_splits), + metadata=metadata, + ) + if is_capturing: + # CUDA graphs retain raw pointers, not Python tensor owners. + # Keep every capture's metadata alive for the graph lifetime. + self._captured_scheduler_plans.append(plan) + else: + self._scheduler_plan = plan + return metadata + + def run_varlen( + self, + q_rope: torch.Tensor, + q_latent: torch.Tensor, + rope_cache: torch.Tensor, + latent_cache: torch.Tensor, + page_table: torch.Tensor, + request_indices: torch.Tensor, + context_lens: torch.Tensor, + output: torch.Tensor, + *, + cu_seqlens_q: torch.Tensor, + max_seqlen_q: int, + num_splits: int | None = None, + validation_scope: object | None = None, + ) -> torch.Tensor: + split_count = self.num_splits if num_splits is None else int(num_splits) + if split_count < 0: + raise ValueError( + f"FA3 num_splits must be non-negative, got {split_count}." + ) + scheduler_metadata = self._scheduler_metadata( + q_rope, + rope_cache, + page_table, + context_lens, + cu_seqlens_q, + headdim_v=int(q_latent.shape[-1]), + max_seqlen_q=int(max_seqlen_q), + num_splits=split_count, + validation_scope=validation_scope, + ) + args: list[object] = [ + q_rope, + rope_cache.unsqueeze(1), + latent_cache.unsqueeze(1), + None, + None, + q_latent, + output, + cu_seqlens_q, + None, + None, + None, + context_lens, + int(max_seqlen_q), + None, + page_table, + request_indices, + None, + None, + None, + None, + None, + None, + None, + self.softmax_scale, + True, + -1, + -1, + ] + if self._has_attention_chunk: + args.append(0) + args.extend( + (0.0, True, scheduler_metadata, split_count, None, 0, None) + ) + if self._has_04_arguments: + args.extend((None, False)) + result: Sequence[torch.Tensor] = self._op(*args) + if not result or result[0].data_ptr() != output.data_ptr(): + raise RuntimeError("sgl-kernel FA3 did not write to the supplied output") + return output + + def run_explicit_varlen( + self, + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + page_table: torch.Tensor, + request_indices: torch.Tensor, + context_lens: torch.Tensor, + output: torch.Tensor, + *, + cu_seqlens_q: torch.Tensor, + max_seqlen_q: int, + validation_scope: object | None = None, + ) -> torch.Tensor: + """Run causal varlen attention over page-size-one explicit KV.""" + + scheduler_metadata = self._scheduler_metadata( + q, + k_cache, + page_table, + context_lens, + cu_seqlens_q, + headdim_v=int(v_cache.shape[-1]), + max_seqlen_q=int(max_seqlen_q), + num_splits=self.num_splits, + validation_scope=validation_scope, + ) + + args: list[object] = [ + q, + k_cache.unsqueeze(1), + v_cache.unsqueeze(1), + None, + None, + None, + output, + cu_seqlens_q, + None, + None, + None, + context_lens, + int(max_seqlen_q), + None, + page_table, + request_indices, + None, + None, + None, + None, + None, + None, + None, + self.softmax_scale, + True, + -1, + -1, + ] + if self._has_attention_chunk: + args.append(0) + args.extend( + (0.0, True, scheduler_metadata, self.num_splits, None, 0, None) + ) + if self._has_04_arguments: + args.extend((None, False)) + result: Sequence[torch.Tensor] = self._op(*args) + if not result or result[0].data_ptr() != output.data_ptr(): + raise RuntimeError("sgl-kernel FA3 did not write to the supplied output") + return output + + +__all__ = ["SglFa3DecodeKernel", "sgl_fa3_support"] diff --git a/src/sparsevllm/operators/sgl_moe.py b/src/sparsevllm/operators/sgl_moe.py new file mode 100644 index 00000000..5842d6a1 --- /dev/null +++ b/src/sparsevllm/operators/sgl_moe.py @@ -0,0 +1,248 @@ +from __future__ import annotations + +import importlib.metadata +import importlib.util +import re + +import torch +import triton + +from sparsevllm.triton_kernel.moe import MoeAlignment + + +_SGL_KERNEL_DISTRIBUTIONS = ("sglang-kernel", "sgl-kernel") +_MIN_VALIDATED_VERSIONS = { + (0, 3): (0, 3, 14), + (0, 4): (0, 4, 5), +} + + +def _parse_version(version: str) -> tuple[int, int, int] | None: + parts = str(version).split(".") + if len(parts) < 3: + return None + parsed: list[int] = [] + for part in parts[:3]: + match = re.match(r"(\d+)", part) + if match is None: + return None + parsed.append(int(match.group(1))) + return tuple(parsed) # type: ignore[return-value] + + +def _installed_sgl_kernel_version() -> tuple[str, tuple[int, int, int]] | None: + for distribution in _SGL_KERNEL_DISTRIBUTIONS: + try: + version = importlib.metadata.version(distribution) + except importlib.metadata.PackageNotFoundError: + continue + parsed = _parse_version(version) + if parsed is None: + return version, (-1, -1, -1) + return version, parsed + return None + + +def _sgl_moe_support(feature: str) -> tuple[bool, str]: + try: + module_spec = importlib.util.find_spec("sgl_kernel") + except (ImportError, ValueError): + module_spec = None + if module_spec is None: + return False, "sgl-kernel is not installed" + installed = _installed_sgl_kernel_version() + if installed is None: + return False, "sgl-kernel package metadata is unavailable" + version, parsed = installed + minimum = _MIN_VALIDATED_VERSIONS.get(parsed[:2]) + if minimum is None: + return ( + False, + "requires validated sgl-kernel 0.3.x or sglang-kernel 0.4.x " + f"API, got {version}", + ) + if parsed < minimum: + package = "sglang-kernel" if parsed[:2] == (0, 4) else "sgl-kernel" + required = ".".join(str(value) for value in minimum) + return False, f"requires {package} >= {required}, got {version}" + return True, f"SGL kernel {version} {feature} is available" + + +def sgl_fused_moe_gate_support() -> tuple[bool, str]: + """Check the SGL fused-gate API used by the GLM router.""" + + return _sgl_moe_support("MoE gate") + + +def sgl_moe_alignment_support() -> tuple[bool, str]: + """Check the SGL expert-alignment API used by the Triton MoE provider.""" + + return _sgl_moe_support("MoE alignment") + + +class SglGlmFusedMoeGate: + """One-kernel sigmoid, biased top-k, normalization, and route scaling.""" + + def __init__(self, *, num_experts: int, top_k: int) -> None: + supported, reason = sgl_fused_moe_gate_support() + if not supported: + raise RuntimeError(reason) + if (int(num_experts), int(top_k)) != (64, 4): + raise ValueError( + "SGL GLM fused gate is validated only for 64 experts and " + f"top-k 4, got {num_experts} and {top_k}." + ) + installed = _installed_sgl_kernel_version() + if installed is None: + raise RuntimeError("sgl-kernel package metadata is unavailable") + _, parsed = installed + self._api_series = parsed[:2] + if self._api_series == (0, 3): + from sgl_kernel import moe_fused_gate + + self._op = moe_fused_gate + else: + from sgl_kernel import topk_sigmoid + + self._op = topk_sigmoid + + def __call__( + self, + router_logits: torch.Tensor, + correction_bias: torch.Tensor, + *, + top_k: int, + routed_scaling_factor: float, + ) -> tuple[torch.Tensor, torch.Tensor]: + if tuple(router_logits.shape[1:]) != (64,) or int(top_k) != 4: + raise ValueError( + "SGL GLM fused gate expects logits [tokens, 64] and top-k 4, " + f"got {tuple(router_logits.shape)} and {top_k}." + ) + if router_logits.dtype != torch.float32: + raise TypeError( + "SGL GLM fused gate requires FP32 logits, got " + f"{router_logits.dtype}." + ) + if tuple(correction_bias.shape) != (64,): + raise ValueError( + "SGL GLM fused gate expects correction_bias [64], got " + f"{tuple(correction_bias.shape)}." + ) + + # The 0.3.x kernel supports at most 32 experts per group. Splitting + # GLM's single 64-expert group into two virtual groups and selecting + # both preserves the original global top-k candidate set. + if self._api_series == (0, 3): + weights, ids = self._op( + router_logits, + correction_bias, + 2, + 2, + 4, + 0, + float(routed_scaling_factor), + True, + ) + else: + weights = torch.empty( + (int(router_logits.shape[0]), 4), + dtype=torch.float32, + device=router_logits.device, + ) + ids = torch.empty( + (int(router_logits.shape[0]), 4), + dtype=torch.int32, + device=router_logits.device, + ) + self._op(weights, ids, router_logits, True, correction_bias) + weights.mul_(float(routed_scaling_factor)) + return weights, ids + + +def sgl_moe_align_block_size( + topk_ids: torch.Tensor, + *, + block_size: int, + num_experts: int, + local_expert_start: int, + local_expert_end: int, +) -> MoeAlignment: + """Build a full-expert MoE assignment with the SGL CUDA kernel.""" + + if (int(local_expert_start), int(local_expert_end)) != (0, int(num_experts)): + raise ValueError("SGL MoE alignment currently requires EP size 1.") + num_assignments = int(topk_ids.numel()) + max_num_tokens_padded = triton.cdiv( + num_assignments + int(num_experts) * (int(block_size) - 1), + int(block_size), + ) * int(block_size) + sorted_token_ids = torch.empty( + max_num_tokens_padded, + dtype=torch.int32, + device=topk_ids.device, + ) + expert_ids = torch.empty( + max_num_tokens_padded // int(block_size), + dtype=torch.int32, + device=topk_ids.device, + ) + num_tokens_post_padded = torch.empty( + 1, + dtype=torch.int32, + device=topk_ids.device, + ) + cumsum_buffer = torch.empty( + int(num_experts) + 1, + dtype=torch.int32, + device=topk_ids.device, + ) + supported, reason = sgl_moe_alignment_support() + if not supported: + raise RuntimeError(reason) + installed = _installed_sgl_kernel_version() + if installed is None: + raise RuntimeError("sgl-kernel package metadata is unavailable") + _, parsed = installed + from sgl_kernel import moe_align_block_size + + # Both validated APIs iterate to num_experts - 1. The extra empty logical + # expert makes the complete [0, num_experts) range participate. + if parsed[:2] == (0, 3): + moe_align_block_size( + topk_ids, + int(num_experts) + 1, + int(block_size), + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + cumsum_buffer, + True, + ) + else: + moe_align_block_size( + topk_ids, + int(num_experts) + 1, + int(block_size), + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + cumsum_buffer, + True, + False, + ) + return MoeAlignment( + sorted_token_ids=sorted_token_ids, + expert_ids=expert_ids, + num_tokens_post_padded=num_tokens_post_padded, + block_size=int(block_size), + naive=False, + ) + + +__all__ = [ + "SglGlmFusedMoeGate", + "sgl_fused_moe_gate_support", + "sgl_moe_alignment_support", + "sgl_moe_align_block_size", +] diff --git a/src/sparsevllm/triton_kernel/mla/__init__.py b/src/sparsevllm/triton_kernel/mla/__init__.py index b38d8245..0106939a 100644 --- a/src/sparsevllm/triton_kernel/mla/__init__.py +++ b/src/sparsevllm/triton_kernel/mla/__init__.py @@ -3,6 +3,7 @@ from .copy_latent import copy_latent_to_cache, validate_copy_slot_mapping from .decode_schedule import ( DEFAULT_GLM_MLA_DECODE_CONFIG, + GLM_MLA_MAX_WORKSPACE_CONFIG, GLM_MLA_SOFTMAX_SCALE, MlaDecodeLaunchConfig, MlaDecodeWorkspace, @@ -10,6 +11,7 @@ prepare_mla_decode_schedule, required_workspace_blocks, run_mla_decode, + select_glm_mla_decode_config, validate_mla_decode_metadata, ) from .decode_stage1 import MLA_LATENT_DIM, MLA_ROPE_DIM, decode_stage1 @@ -21,6 +23,7 @@ __all__ = [ "DEFAULT_GLM_MLA_DECODE_CONFIG", + "GLM_MLA_MAX_WORKSPACE_CONFIG", "GLM_MLA_SOFTMAX_SCALE", "MLA_LATENT_DIM", "MLA_ROPE_DIM", @@ -34,6 +37,7 @@ "prepare_mla_decode_schedule", "required_workspace_blocks", "run_mla_decode", + "select_glm_mla_decode_config", "validate_copy_slot_mapping", "validate_gather_metadata", "validate_mla_decode_metadata", diff --git a/src/sparsevllm/triton_kernel/mla/decode_schedule.py b/src/sparsevllm/triton_kernel/mla/decode_schedule.py index c11f6b52..96086820 100644 --- a/src/sparsevllm/triton_kernel/mla/decode_schedule.py +++ b/src/sparsevllm/triton_kernel/mla/decode_schedule.py @@ -58,6 +58,81 @@ def __post_init__(self) -> None: DEFAULT_GLM_MLA_DECODE_CONFIG = MlaDecodeLaunchConfig() +# Measured on NVIDIA H100 80GB HBM3 with GLM-4.7-Flash TP2. Keep this +# table deliberately narrow: other TP layouts retain the correctness-first +# default until they have their own matched-shape measurements. +_GLM_MLA_TP2_SMALL_BATCH_CONFIG = MlaDecodeLaunchConfig( + program_count=256, + blocks_per_program=4, + block_n=32, + block_q_heads=16, + stage1_num_warps=8, + stage1_pipeline_stages=6, + stage2_num_warps=4, + stage2_pipeline_stages=1, +) +_GLM_MLA_TP2_MEDIUM_BATCH_CONFIG = MlaDecodeLaunchConfig( + program_count=264, + blocks_per_program=2, + block_n=32, + block_q_heads=8, + stage1_num_warps=8, + stage1_pipeline_stages=4, + stage2_num_warps=4, + stage2_pipeline_stages=1, +) +_GLM_MLA_TP2_SHORT_CONTEXT_CONFIG = MlaDecodeLaunchConfig( + program_count=128, + blocks_per_program=8, + block_n=32, + block_q_heads=8, + stage1_num_warps=8, + stage1_pipeline_stages=4, + stage2_num_warps=4, + stage2_pipeline_stages=1, +) +_GLM_MLA_TP2_LARGE_BATCH_CONFIG = MlaDecodeLaunchConfig( + program_count=256, + blocks_per_program=8, + block_n=32, + block_q_heads=8, + stage1_num_warps=8, + stage1_pipeline_stages=4, + stage2_num_warps=4, + stage2_pipeline_stages=1, +) + +# One caller-owned allocation accommodates every measured TP2 schedule. +GLM_MLA_MAX_WORKSPACE_CONFIG = MlaDecodeLaunchConfig( + program_count=264, + blocks_per_program=8, +) + + +def select_glm_mla_decode_config( + *, + batch_size: int, + max_context_len: int, + local_q_heads: int, +) -> MlaDecodeLaunchConfig: + """Select a graph-stable launch config from static decode dimensions.""" + + if batch_size <= 0: + raise ValueError("batch_size must be positive") + if max_context_len <= 0: + raise ValueError("max_context_len must be positive") + if local_q_heads <= 0: + raise ValueError("local_q_heads must be positive") + if local_q_heads != 10: + return DEFAULT_GLM_MLA_DECODE_CONFIG + if batch_size <= 1: + return _GLM_MLA_TP2_SMALL_BATCH_CONFIG + if batch_size <= 8: + return _GLM_MLA_TP2_MEDIUM_BATCH_CONFIG + if max_context_len <= 1024: + return _GLM_MLA_TP2_SHORT_CONTEXT_CONFIG + return _GLM_MLA_TP2_LARGE_BATCH_CONFIG + @dataclass(frozen=True, slots=True) class MlaDecodeWorkspace: diff --git a/src/sparsevllm/triton_kernel/moe.py b/src/sparsevllm/triton_kernel/moe.py index cc0e2088..a2b93303 100644 --- a/src/sparsevllm/triton_kernel/moe.py +++ b/src/sparsevllm/triton_kernel/moe.py @@ -1,5 +1,6 @@ from __future__ import annotations +from collections.abc import Callable from dataclasses import dataclass import torch @@ -27,6 +28,89 @@ class MoeAlignment: naive: bool +@triton.jit +def _append_shared_expert_route_kernel( + input_ids_ptr, + input_weights_ptr, + output_ids_ptr, + output_weights_ptr, + shared_expert_id: tl.constexpr, + INPUT_TOP_K: tl.constexpr, + OUTPUT_TOP_K: tl.constexpr, + BLOCK_TOP_K: tl.constexpr, +): + token_id = tl.program_id(0) + route_offsets = tl.arange(0, BLOCK_TOP_K) + input_offsets = token_id * INPUT_TOP_K + route_offsets + output_offsets = token_id * OUTPUT_TOP_K + route_offsets + routed = route_offsets < INPUT_TOP_K + route_ids = tl.load(input_ids_ptr + input_offsets, mask=routed, other=0) + route_weights = tl.load( + input_weights_ptr + input_offsets, + mask=routed, + other=0.0, + ) + route_ids = tl.where(routed, route_ids, shared_expert_id) + route_weights = tl.where(routed, route_weights, 1.0) + output_mask = route_offsets < OUTPUT_TOP_K + tl.store(output_ids_ptr + output_offsets, route_ids, mask=output_mask) + tl.store(output_weights_ptr + output_offsets, route_weights, mask=output_mask) + + +def append_shared_expert_route( + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + *, + shared_expert_id: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Append one always-active shared expert in one graph-safe kernel.""" + + if not topk_ids.is_cuda or not topk_weights.is_cuda: + raise ValueError("Shared-expert route packing requires CUDA tensors.") + if topk_ids.device != topk_weights.device: + raise ValueError("Shared-expert route tensors must share one device.") + if topk_ids.ndim != 2 or topk_weights.shape != topk_ids.shape: + raise ValueError( + "Shared-expert route tensors must share [tokens, top_k] shape." + ) + if topk_ids.dtype not in (torch.int32, torch.int64): + raise TypeError("Shared-expert route IDs must use int32 or int64.") + if topk_weights.dtype not in (*_SUPPORTED_DTYPES, torch.float32): + raise TypeError("Shared-expert route weights must use BF16, FP16, or FP32.") + if not topk_ids.is_contiguous() or not topk_weights.is_contiguous(): + raise ValueError("Shared-expert route tensors must be contiguous.") + if topk_ids.shape[0] <= 0 or topk_ids.shape[1] <= 0: + raise ValueError("Shared-expert route tensors must be non-empty.") + shared_expert_id = int(shared_expert_id) + if shared_expert_id < 0: + raise ValueError("shared_expert_id must be non-negative.") + + tokens, input_top_k = map(int, topk_ids.shape) + output_top_k = input_top_k + 1 + output_ids = torch.empty( + (tokens, output_top_k), + dtype=topk_ids.dtype, + device=topk_ids.device, + ) + output_weights = torch.empty( + (tokens, output_top_k), + dtype=topk_weights.dtype, + device=topk_weights.device, + ) + _append_shared_expert_route_kernel[(tokens,)]( + topk_ids, + topk_weights, + output_ids, + output_weights, + shared_expert_id=shared_expert_id, + INPUT_TOP_K=input_top_k, + OUTPUT_TOP_K=output_top_k, + BLOCK_TOP_K=triton.next_power_of_2(output_top_k), + num_warps=1, + ) + return output_ids, output_weights + + @triton.jit( do_not_specialize=[ "num_assignments", @@ -1084,6 +1168,7 @@ def fused_moe( local_expert_start: int, output_dtype: torch.dtype | None = None, _fuse_gate_up_swiglu: bool = False, + alignment_impl: Callable[..., MoeAlignment] | None = None, ) -> torch.Tensor: """Run unquantized routed experts with a generic Triton MoE pipeline. @@ -1126,7 +1211,9 @@ def fused_moe( device_name=device_name, device_capability=capability, ).as_triton_kwargs() - alignment = _prepare_expert_assignment( + if alignment_impl is None: + alignment_impl = _prepare_expert_assignment + alignment = alignment_impl( topk_ids, block_size=w13_config["BLOCK_SIZE_M"], num_experts=num_experts, diff --git a/src/sparsevllm/triton_kernel/moe_config.py b/src/sparsevllm/triton_kernel/moe_config.py index d19837ca..0fe696d1 100644 --- a/src/sparsevllm/triton_kernel/moe_config.py +++ b/src/sparsevllm/triton_kernel/moe_config.py @@ -96,6 +96,49 @@ def _heuristic_config( _I = MoeGemmConfig(16, 64, 64, 8, 4, 4) _J = MoeGemmConfig(16, 64, 64, 8, 4, 2) _K = MoeGemmConfig(16, 128, 64, 8, 4, 2) +_GLM_DECODE_32 = MoeGemmConfig(16, 64, 128, 16, 4, 3) +_GLM_DECODE_64 = MoeGemmConfig(64, 128, 64, 8, 8, 3) +_GLM_MID_BATCH = MoeGemmConfig(64, 128, 64, 1, 8, 3) +_GLM_LARGE_BATCH = MoeGemmConfig(128, 128, 64, 1, 8, 3) + + +def _glm_h100_tp2_config( + shape: MoeGemmShape, + *, + num_tokens: int, + stage: str, +) -> MoeGemmConfig | None: + """Return measured BF16 configs for the GLM TP2 expert shape.""" + + profiled_shapes = { + MoeGemmShape( + "NVIDIA H100 80GB HBM3", + (9, 0), + torch.bfloat16, + 4, + 64, + 2048, + 768, + ), + MoeGemmShape( + "NVIDIA H100 80GB HBM3", + (9, 0), + torch.bfloat16, + 5, + 65, + 2048, + 768, + ), + } + if stage not in {"w13", "w2"} or shape not in profiled_shapes: + return None + if num_tokens <= 32: + return _GLM_DECODE_32 + if num_tokens <= 64: + return _GLM_DECODE_64 + if num_tokens <= 512: + return _GLM_MID_BATCH + return _GLM_LARGE_BATCH def _stage_table( @@ -391,6 +434,13 @@ def _resolve_moe_gemm_config( hidden_size=hidden_size, intermediate_size=intermediate_size, ) + glm_config = _glm_h100_tp2_config( + shape, + num_tokens=num_tokens, + stage=stage, + ) + if glm_config is not None: + return glm_config table = _TUNED_CONFIGS.get(shape) if stage == "gate_up_swiglu": fused_table = _TUNED_GATE_UP_SWIGLU_CONFIGS.get(shape) diff --git a/src/sparsevllm/triton_kernel/silu_and_mul.py b/src/sparsevllm/triton_kernel/silu_and_mul.py index e3e7ef51..b174fec6 100644 --- a/src/sparsevllm/triton_kernel/silu_and_mul.py +++ b/src/sparsevllm/triton_kernel/silu_and_mul.py @@ -58,6 +58,12 @@ def _silu_and_mul_kernel( ) +def _resolve_silu_launch_config(size_m: int) -> tuple[int, int, int | None]: + if int(size_m) <= 256: + return 32, 128, 4 + return 128, 128, None + + def silu_and_mul_fwd(input, *, gate_up_order: str = "gate_up"): if gate_up_order not in {"gate_up", "up_gate"}: raise ValueError( @@ -70,12 +76,12 @@ def silu_and_mul_fwd(input, *, gate_up_order: str = "gate_up"): stride_output_n = input.stride(1) size_m = input.shape[0] size_n = input.shape[-1] // 2 - BLOCK_M = 128 - BLOCK_N = 128 + BLOCK_M, BLOCK_N, num_warps = _resolve_silu_launch_config(size_m) grid = ( triton.cdiv(size_m, BLOCK_M), triton.cdiv(size_n, BLOCK_N), ) + launch_kwargs = {} if num_warps is None else {"num_warps": num_warps} _silu_and_mul_kernel[grid]( input, stride_input_m, @@ -87,6 +93,7 @@ def silu_and_mul_fwd(input, *, gate_up_order: str = "gate_up"): GATE_FIRST=gate_up_order == "gate_up", BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, + **launch_kwargs, ) return input[:, 0 : (input.shape[-1] // 2)] diff --git a/tests/test_glm4_moe_lite.py b/tests/test_glm4_moe_lite.py index 6ba0e5d5..9019ae1c 100644 --- a/tests/test_glm4_moe_lite.py +++ b/tests/test_glm4_moe_lite.py @@ -34,6 +34,7 @@ from sparsevllm.models.qwen3 import Qwen3MLP from sparsevllm.operators.mla_attention import MlaAttentionOpSpec from sparsevllm.operators.moe import TritonMoeProvider +from sparsevllm.platforms import device_runtime def _config(**overrides) -> Glm4MoeLiteConfig: @@ -225,6 +226,45 @@ def test_glm_tp_projection_slices_follow_local_heads() -> None: assert torch.equal(attention.o_proj.weight, o_source[:, 2560:3840]) +def test_glm_qkv_a_projection_loads_and_executes_as_one_gemm() -> None: + config = _config() + context = _tp_context() + with _construction_context(context): + attention = Glm4MoeLiteAttention( + config, + _fake_mla(), + projection_chunk_size=config.mlp_chunk_size, + ) + + torch.manual_seed(23) + q_weight = torch.randn(config.q_lora_rank, config.hidden_size) + kv_output_size = config.kv_lora_rank + config.qk_rope_head_dim + kv_weight = torch.randn(kv_output_size, config.hidden_size) + projection = attention.fused_qkv_a_proj + projection.weight_loader(projection.weight, q_weight, 0) + projection.weight_loader(projection.weight, kv_weight, 1) + + q_weight = q_weight.to(dtype=projection.weight.dtype) + kv_weight = kv_weight.to(dtype=projection.weight.dtype) + hidden_states = torch.randn( + 5, + config.hidden_size, + dtype=projection.weight.dtype, + ) + actual_q, actual_kv = projection(hidden_states).split( + [config.q_lora_rank, kv_output_size], + dim=-1, + ) + torch.testing.assert_close(actual_q, F.linear(hidden_states, q_weight)) + torch.testing.assert_close(actual_kv, F.linear(hidden_states, kv_weight)) + assert Glm4MoeLiteForCausalLM.packed_modules_mapping[ + "self_attn.q_a_proj" + ] == ("self_attn.fused_qkv_a_proj", 0) + assert Glm4MoeLiteForCausalLM.packed_modules_mapping[ + "self_attn.kv_a_proj_with_mqa" + ] == ("self_attn.fused_qkv_a_proj", 1) + + def test_glm_decode_absorption_and_value_reconstruction_match_linear_algebra() -> None: config = _config() context = _tp_context() @@ -395,7 +435,11 @@ def test_glm_sparse_moe_reduces_pure_ep_over_world(ep_size: int) -> None: assert all_reduce.call_args.kwargs["group"] is context.world.process_group -def test_glm_sparse_moe_reduces_pure_tp_over_world() -> None: +@pytest.mark.parametrize(("is_prefill", "expected_reductions"), [(False, 1), (True, 1)]) +def test_glm_sparse_moe_reduces_pure_tp_over_world( + is_prefill: bool, + expected_reductions: int, +) -> None: moe_tp_process_group = object() context = ParallelContext( world=ParallelGroup(moe_tp_process_group, (0, 1), 0, 2), @@ -412,12 +456,21 @@ def test_glm_sparse_moe_reduces_pure_tp_over_world() -> None: block._routed_chunk = lambda hidden_states: hidden_states.clone() hidden_states = torch.arange(8, dtype=torch.float32).reshape(2, 4) - with patch.object(dist, "all_reduce") as all_reduce: + with ( + patch.object(dist, "all_reduce") as all_reduce, + patch( + "sparsevllm.models.glm4_moe_lite.get_context", + return_value=SimpleNamespace(is_prefill=is_prefill), + ), + ): output = block(hidden_states) torch.testing.assert_close(output, hidden_states * 2) - all_reduce.assert_called_once() - assert all_reduce.call_args.kwargs["group"] is context.world.process_group + assert all_reduce.call_count == expected_reductions + assert all( + call.kwargs["group"] is context.world.process_group + for call in all_reduce.call_args_list + ) def test_glm_sparse_moe_reduces_hybrid_tp_ep_shards_over_outer_world() -> None: @@ -475,10 +528,18 @@ def forward(self, hidden_states, topk_ids, topk_weights): block.shared_experts = nn.Identity() hidden_states = torch.arange(8, dtype=torch.float32).reshape(2, 4) - with patch.dict(os.environ, {"SPARSEVLLM_DEBUG_MOE": "1"}): + with ( + patch.dict(os.environ, {"SPARSEVLLM_DEBUG_MOE": "1"}), + patch.object( + device_runtime, + "is_stream_capturing", + return_value=True, + ), + ): output = block(hidden_states) torch.testing.assert_close(output, hidden_states * 3) + assert isinstance(block.debug_last_local_hit_count, torch.Tensor) torch.testing.assert_close(block.debug_last_local_output, hidden_states * 2) torch.testing.assert_close(block.debug_last_routed_output, hidden_states * 2) torch.testing.assert_close(block.debug_last_output, output) diff --git a/tests/test_glm_cuda_graph.py b/tests/test_glm_cuda_graph.py index 8414b23b..85ac0f5e 100644 --- a/tests/test_glm_cuda_graph.py +++ b/tests/test_glm_cuda_graph.py @@ -249,6 +249,7 @@ def run_model( is_long_text_batch=lambda seqs, is_prefill: False, method="", rank=0, + parallel_context=parallel_context, capture_sizes=[1], context_sizes=[128], ) @@ -558,6 +559,7 @@ def run_model( is_long_text_batch=lambda seqs, is_prefill: False, method="", rank=0, + parallel_context=parallel_context, capture_sizes=[1], context_sizes=[128], ) @@ -1117,6 +1119,7 @@ def run_model( is_long_text_batch=lambda seqs, is_prefill: True, method=method, rank=0, + parallel_context=parallel_context, capture_sizes=[1], context_sizes=[16], ) diff --git a/tests/test_mla_attention_layer.py b/tests/test_mla_attention_layer.py index 2647738b..49424c95 100644 --- a/tests/test_mla_attention_layer.py +++ b/tests/test_mla_attention_layer.py @@ -40,6 +40,30 @@ def __init__( self.max_batch_size = int(max_batch_size) +class _ExplicitPrefillProvider(_TestProvider): + supports_explicit_prefill = True + + def run_explicit_prefill( + self, + q, + view, + output, + *, + cu_seqlens_q, + max_seqlen_q, + validation_scope=None, + ): + self.prefill_call = { + "q": q, + "view": view, + "cu_seqlens_q": cu_seqlens_q, + "max_seqlen_q": max_seqlen_q, + "validation_scope": validation_scope, + } + output.copy_(q) + return output + + def _spec(tp_size: int = 4) -> MlaAttentionOpSpec: return MlaAttentionOpSpec( num_q_heads=20, @@ -367,6 +391,45 @@ def test_mla_prefill_exposes_expanded_explicit_score_view() -> None: assert score_view.meta.max_context_len == history.max_context_len +def test_mla_prefill_dispatches_explicit_provider_with_shared_cu_seqlens() -> None: + spec = _spec() + provider = _ExplicitPrefillProvider(spec, device="cpu", max_batch_size=4) + attention = MLAAttention( + spec=spec, + provider=provider, + prefill_workspace_bytes=64 * 1024 * 1024, + hidden_size=64, + projection_chunk_size=8, + ) + view = _view( + torch.empty(2, 1, 512, dtype=torch.bfloat16), + torch.empty(2, 1, 64, dtype=torch.bfloat16), + torch.tensor([[0, 1]], dtype=torch.int32), + torch.tensor([0], dtype=torch.int32), + torch.tensor([2], dtype=torch.int32), + ) + with patch("sparsevllm.layers.mla_attention.gather_latent_history"): + history = attention.prepare_prefill_history(view, query_tokens=2) + workset = _expand_history(attention, history) + q = torch.randn(2, 5, 256, dtype=torch.bfloat16) + cu_seqlens_q = torch.tensor([0, 2], dtype=torch.int32) + set_context(True, cu_seqlens_q=cu_seqlens_q) + try: + output = attention.run_prefill( + q, + workset, + b_start_loc=cu_seqlens_q[:-1], + chunk_lens=cu_seqlens_q[1:] - cu_seqlens_q[:-1], + ) + finally: + reset_context() + + torch.testing.assert_close(output, q) + assert provider.prefill_call["view"].payload.k_cache is workset.expanded_k + assert provider.prefill_call["cu_seqlens_q"] is cu_seqlens_q + assert provider.prefill_call["max_seqlen_q"] == 2 + + def test_mla_materializes_actual_keys_for_permuted_slots() -> None: torch.manual_seed(47) attention = _attention(tp_size=4) diff --git a/tests/test_mla_attention_operator.py b/tests/test_mla_attention_operator.py index c27f32ab..87040b2d 100644 --- a/tests/test_mla_attention_operator.py +++ b/tests/test_mla_attention_operator.py @@ -14,11 +14,15 @@ from sparsevllm.operators.mla_attention import ( MLA_ATTENTION_REGISTRY, MlaAttentionOpSpec, + MlaSglFa3Provider, MlaTritonProvider, ) from sparsevllm.operators.registry import OpResolver from sparsevllm.platforms import DeviceCaps, PlatformEnum -from sparsevllm.triton_kernel.mla import MlaDecodeWorkspace +from sparsevllm.triton_kernel.mla import ( + GLM_MLA_MAX_WORKSPACE_CONFIG, + MlaDecodeWorkspace, +) def _spec(**overrides) -> MlaAttentionOpSpec: @@ -93,10 +97,16 @@ def test_mla_resolver_selects_h100_triton_provider() -> None: spec = _spec() workspace = _cpu_workspace(batch_size=8, head_count=spec.local_q_heads) - with patch( - "sparsevllm.operators.mla_attention.allocate_mla_decode_workspace", - return_value=workspace, - ) as allocate: + with ( + patch( + "sparsevllm.operators.mla_attention.sgl_fa3_support", + return_value=(False, "sgl-kernel is not installed"), + ), + patch( + "sparsevllm.operators.mla_attention.allocate_mla_decode_workspace", + return_value=workspace, + ) as allocate, + ): resolved = OpResolver(MLA_ATTENTION_REGISTRY).resolve( spec, _h100_caps(), @@ -106,15 +116,44 @@ def test_mla_resolver_selects_h100_triton_provider() -> None: ) assert isinstance(resolved.provider, MlaTritonProvider) - assert resolved.rejected == () + assert resolved.rejected == ( + ("sgl_fa3_h100", "sgl-kernel is not installed"), + ) allocate.assert_called_once_with( batch_size=8, head_count=5, device=torch.device("cpu"), - config=resolved.provider.launch_config, + config=GLM_MLA_MAX_WORKSPACE_CONFIG, ) +def test_mla_resolver_prefers_sgl_fa3_when_available() -> None: + spec = _spec() + workspace = _cpu_workspace(batch_size=8, head_count=spec.local_q_heads) + + with ( + patch( + "sparsevllm.operators.mla_attention.sgl_fa3_support", + return_value=(True, "validated test API"), + ), + patch( + "sparsevllm.operators.mla_attention.allocate_mla_decode_workspace", + return_value=workspace, + ), + patch("sparsevllm.operators.mla_attention.SglFa3DecodeKernel"), + ): + resolved = OpResolver(MLA_ATTENTION_REGISTRY).resolve( + spec, + _h100_caps(), + op_spec=spec, + device="cpu", + max_batch_size=8, + ) + + assert type(resolved.provider) is MlaSglFa3Provider + assert resolved.rejected == () + + def test_mla_resolver_accepts_cuda_graph_after_capture_gate() -> None: spec = _spec(cuda_graph=True) workspace = _cpu_workspace(batch_size=8, head_count=spec.local_q_heads) diff --git a/tests/test_mla_kernels.py b/tests/test_mla_kernels.py index fb5965c7..8f9de29b 100644 --- a/tests/test_mla_kernels.py +++ b/tests/test_mla_kernels.py @@ -18,6 +18,7 @@ gather_latent_history, prepare_mla_decode_schedule, run_mla_decode, + select_glm_mla_decode_config, ) @@ -28,6 +29,35 @@ DECODE_CONTEXTS = (1, 31, 32, 33, 127, 128, 129, 255, 256, 257, 1024, 4096) +def test_select_glm_mla_decode_config_uses_measured_tp2_shapes() -> None: + small = select_glm_mla_decode_config( + batch_size=1, max_context_len=4096, local_q_heads=10 + ) + medium = select_glm_mla_decode_config( + batch_size=8, max_context_len=4096, local_q_heads=10 + ) + short = select_glm_mla_decode_config( + batch_size=32, max_context_len=1024, local_q_heads=10 + ) + large = select_glm_mla_decode_config( + batch_size=32, max_context_len=4096, local_q_heads=10 + ) + + assert (small.program_count, small.blocks_per_program) == (256, 4) + assert (medium.program_count, medium.blocks_per_program) == (264, 2) + assert (short.program_count, short.blocks_per_program) == (128, 8) + assert (large.program_count, large.blocks_per_program) == (256, 8) + assert large.block_q_heads == 8 + + +def test_select_glm_mla_decode_config_keeps_unmeasured_tp_default() -> None: + actual = select_glm_mla_decode_config( + batch_size=32, max_context_len=4096, local_q_heads=5 + ) + + assert actual == DEFAULT_GLM_MLA_DECODE_CONFIG + + def _torch_mla_decode( q_latent: torch.Tensor, q_rope: torch.Tensor, diff --git a/tests/test_moe_config.py b/tests/test_moe_config.py index a60cf648..66dbc98b 100644 --- a/tests/test_moe_config.py +++ b/tests/test_moe_config.py @@ -47,6 +47,27 @@ def test_moe_config_rejects_unknown_stage(): resolve_moe_gemm_config(**arguments, stage="w3") +def test_glm_fused_shared_decode_reuses_profiled_tile(): + config = resolve_moe_gemm_config( + dtype=torch.bfloat16, + num_tokens=32, + top_k=5, + num_local_experts=65, + hidden_size=2048, + intermediate_size=768, + stage="w2", + device_name="NVIDIA H100 80GB HBM3", + device_capability=(9, 0), + ) + + assert ( + config.block_m, + config.block_n, + config.block_k, + config.group_m, + ) == (16, 64, 128, 16) + + def test_h100_tp_ep_fused_gate_up_uses_dedicated_profile(): common = dict( dtype=torch.bfloat16, @@ -349,3 +370,37 @@ def test_fp8_routed_unknown_shape_uses_explicit_default(): ) assert (config.block_n, config.block_k, config.swap_ab) == (128, 128, False) + + +@pytest.mark.parametrize( + ("num_tokens", "expected"), + [ + (32, (16, 64, 128, 16)), + (64, (64, 128, 64, 8)), + (512, (64, 128, 64, 1)), + (1024, (128, 128, 64, 1)), + (65536, (128, 128, 64, 1)), + ], +) +def test_glm_h100_tp2_profile_covers_decode_and_large_prefill( + num_tokens, + expected, +): + common = dict( + dtype=torch.bfloat16, + num_tokens=num_tokens, + top_k=4, + num_local_experts=64, + hidden_size=2048, + intermediate_size=768, + device_name="NVIDIA H100 80GB HBM3", + device_capability=(9, 0), + ) + for stage in ("w13", "w2"): + config = resolve_moe_gemm_config(**common, stage=stage) + assert ( + config.block_m, + config.block_n, + config.block_k, + config.group_m, + ) == expected diff --git a/tests/test_operator_providers.py b/tests/test_operator_providers.py index f6357d45..0748300d 100644 --- a/tests/test_operator_providers.py +++ b/tests/test_operator_providers.py @@ -453,6 +453,68 @@ def test_hopper_fused_moe_uses_profiled_tp_ep_shape(): assert resolved.provider.name == "triton_hopper_fused" +def test_glm_tp2_moe_uses_sgl_aligned_provider(): + spec = _moe_spec( + activation_dtype=torch.bfloat16, + weight_dtype=torch.bfloat16, + block_shape=None, + hidden_size=2048, + intermediate_size=768, + num_local_experts=64, + num_experts=64, + top_k=4, + ep_size=1, + tp_size=2, + routing_method="biased_sigmoid", + ) + + with patch( + "sparsevllm.operators.sgl_moe.sgl_moe_alignment_support", + return_value=(True, "available"), + ): + resolved = OpResolver(MOE_REGISTRY).resolve( + spec, + _cuda_caps( + (9, 0), + native_fp8=False, + device_name="NVIDIA H100 80GB HBM3", + ), + ) + + assert resolved.provider.name == "sgl_aligned_triton_glm" + + +def test_glm_tp2_fused_shared_decode_uses_sgl_aligned_provider(): + spec = _moe_spec( + activation_dtype=torch.bfloat16, + weight_dtype=torch.bfloat16, + block_shape=None, + hidden_size=2048, + intermediate_size=768, + num_local_experts=65, + num_experts=65, + top_k=5, + ep_size=1, + tp_size=2, + routing_method="biased_sigmoid", + ) + + with patch( + "sparsevllm.operators.sgl_moe.sgl_moe_alignment_support", + return_value=(True, "available"), + ): + resolved = OpResolver(MOE_REGISTRY).resolve( + spec, + _cuda_caps( + (9, 0), + native_fp8=False, + device_name="NVIDIA H100 80GB HBM3", + ), + ) + + assert resolved.provider.name == "sgl_aligned_triton_glm" + + @pytest.mark.parametrize( ("tp_size", "ep_size", "intermediate_size", "num_local_experts"), [(1, 1, 512, 256), (2, 1, 256, 256), (1, 2, 512, 128)], diff --git a/tests/test_parallel_context.py b/tests/test_parallel_context.py index e46c67a3..ae59b857 100644 --- a/tests/test_parallel_context.py +++ b/tests/test_parallel_context.py @@ -231,6 +231,81 @@ def test_ep_broadcast_rejects_invalid_source_rank(): context.ep_broadcast(torch.tensor([1.0]), src_ep_rank=4) +def test_small_out_of_place_all_reduce_uses_bound_provider(): + ranks = (0, 1) + group = ParallelGroup(object(), ranks, 0, 2) + reduced = torch.tensor([3.0]) + provider = SimpleNamespace( + group=group, + all_reduce=Mock(return_value=reduced), + ) + context = ParallelContext( + world=group, + tensor=group, + expert=ParallelGroup(None, (0,), 0, 1), + data=ParallelGroup(None, (0,), 0, 1), + all_reduce_provider=provider, + ) + tensor = torch.tensor([1.0]) + + with patch.object(dist, "all_reduce") as nccl_all_reduce: + output = context.tp_all_reduce_out_of_place(tensor) + + assert output is reduced + provider.all_reduce.assert_called_once_with(tensor) + nccl_all_reduce.assert_not_called() + + +def test_small_out_of_place_all_reduce_falls_back_to_nccl(): + ranks = (0, 1) + group = ParallelGroup(object(), ranks, 0, 2) + provider = SimpleNamespace(group=group, all_reduce=Mock(return_value=None)) + context = ParallelContext( + world=group, + tensor=group, + expert=ParallelGroup(None, (0,), 0, 1), + data=ParallelGroup(None, (0,), 0, 1), + all_reduce_provider=provider, + ) + tensor = torch.tensor([1.0]) + + with patch.object(dist, "all_reduce") as nccl_all_reduce: + output = context.world_all_reduce_out_of_place(tensor) + + assert output is tensor + nccl_all_reduce.assert_called_once_with( + tensor, + op=dist.ReduceOp.SUM, + group=group.process_group, + ) + + +def test_all_reduce_capture_and_close_delegate_to_provider(): + group = ParallelGroup(None, (0,), 0, 1) + capture = Mock() + capture.return_value.__enter__ = Mock(return_value=None) + capture.return_value.__exit__ = Mock(return_value=False) + provider = SimpleNamespace( + group=group, + capture=capture, + close=Mock(), + ) + context = ParallelContext( + world=group, + tensor=group, + expert=group, + data=group, + all_reduce_provider=provider, + ) + + with context.all_reduce_capture(): + pass + context.close_all_reduce_provider() + + capture.assert_called_once_with() + provider.close.assert_called_once_with() + + def test_qwen3_moe_parallel_config_validation(tmp_path): with patch("sparsevllm.configs.runtime.AutoConfig.from_pretrained", return_value=_hf_config()): config = Config(model=str(tmp_path), expert_parallel_size=4) @@ -406,6 +481,25 @@ def test_parallel_group_uses_bound_all_reduce_provider(): provider.run.assert_called_once_with(input_tensor) +def test_vocab_parallel_embedding_uses_out_of_place_reduction(): + reduced = torch.randn(2, 4) + context = SimpleNamespace( + tp_rank=0, + tp_size=2, + tp_all_reduce_out_of_place=Mock(return_value=reduced), + ) + with patch( + "sparsevllm.layers.embed_head.get_parallel_context", + return_value=context, + ): + embedding = VocabParallelEmbedding(8, 4) + + output = embedding(torch.tensor([0, 5])) + + assert output is reduced + context.tp_all_reduce_out_of_place.assert_called_once() + + def test_cache_kv_heads_depend_on_tp_not_ep(): context = _replicated_ep_context() config = SimpleNamespace( diff --git a/tests/test_qwen3_fp8.py b/tests/test_qwen3_fp8.py index 174c6ecc..02fd59e1 100644 --- a/tests/test_qwen3_fp8.py +++ b/tests/test_qwen3_fp8.py @@ -16,6 +16,7 @@ def _parallel_context(): tp_rank=0, tp_size=1, tp_all_reduce=lambda tensor: tensor, + tp_all_reduce_out_of_place=lambda tensor: tensor, ) diff --git a/tests/test_sgl_fa3.py b/tests/test_sgl_fa3.py new file mode 100644 index 00000000..9b109e5b --- /dev/null +++ b/tests/test_sgl_fa3.py @@ -0,0 +1,343 @@ +from __future__ import annotations + +import importlib.metadata +from unittest.mock import patch + +import pytest +import torch + +from sparsevllm.operators.sgl_fa3 import ( + SglFa3DecodeKernel, + _parse_version, + sgl_fa3_support, +) + + +def test_sgl_fa3_version_parser_accepts_post_release() -> None: + assert _parse_version("0.3.14.post1+cu128") == (0, 3, 14) + + +def test_sgl_fa3_support_rejects_missing_package() -> None: + with patch("importlib.util.find_spec", return_value=None): + assert sgl_fa3_support() == (False, "sgl-kernel is not installed") + + +def test_sgl_fa3_support_rejects_old_package() -> None: + with ( + patch("importlib.util.find_spec", return_value=object()), + patch("importlib.metadata.version", return_value="0.3.13"), + ): + supported, reason = sgl_fa3_support() + + assert not supported + assert ">= 0.3.14" in reason + + +def test_sgl_fa3_support_rejects_old_04_package() -> None: + with ( + patch("importlib.util.find_spec", return_value=object()), + patch("importlib.metadata.version", return_value="0.4.0"), + ): + supported, reason = sgl_fa3_support() + + assert not supported + assert "sglang-kernel >= 0.4.5" in reason + + +def test_sgl_fa3_support_accepts_sglang_kernel_04_package() -> None: + def package_version(distribution: str) -> str: + if distribution == "sglang-kernel": + return "0.4.5" + raise importlib.metadata.PackageNotFoundError(distribution) + + with ( + patch("importlib.util.find_spec", return_value=object()), + patch("importlib.metadata.version", side_effect=package_version), + ): + supported, reason = sgl_fa3_support() + + assert supported + assert "0.4.5" in reason + + +@pytest.mark.skipif( + not torch.cuda.is_available() or not sgl_fa3_support()[0], + reason="CUDA and a validated sgl-kernel are required", +) +def test_sgl_fa3_decode_matches_torch_and_replays_cuda_graph() -> None: + torch.manual_seed(20260807) + device = torch.device("cuda") + batch_size, heads, width = 3, 10, 8 + slots = 4 * width + q_rope = torch.randn( + batch_size, heads, 64, device=device, dtype=torch.bfloat16 + ) + q_latent = torch.randn( + batch_size, heads, 512, device=device, dtype=torch.bfloat16 + ) + rope_cache = torch.randn( + slots, 1, 64, device=device, dtype=torch.bfloat16 + ) + latent_cache = torch.randn( + slots, 1, 512, device=device, dtype=torch.bfloat16 + ) + page_table = torch.arange( + slots, device=device, dtype=torch.int32 + ).view(4, width) + request_indices = torch.tensor( + [2, 0, -1], device=device, dtype=torch.int32 + ) + context_lens = torch.tensor( + [7, 5, 0], device=device, dtype=torch.int32 + ) + output = torch.empty_like(q_latent) + kernel = SglFa3DecodeKernel( + device=device, + max_batch_size=batch_size, + softmax_scale=256**-0.5, + ) + + validation_scope = object() + scheduler_op = kernel._scheduler_op + scheduler_call_count = 0 + if scheduler_op is not None: + + def counted_scheduler_op(*args, **kwargs): + nonlocal scheduler_call_count + scheduler_call_count += 1 + return scheduler_op(*args, **kwargs) + + kernel._scheduler_op = counted_scheduler_op + actual = kernel( + q_rope, + q_latent, + rope_cache, + latent_cache, + page_table, + request_indices, + context_lens, + output, + validation_scope=validation_scope, + ) + kernel( + q_rope, + q_latent, + rope_cache, + latent_cache, + page_table, + request_indices, + context_lens, + output, + validation_scope=validation_scope, + ) + assert scheduler_call_count == int(scheduler_op is not None) + kernel( + q_rope, + q_latent, + rope_cache, + latent_cache, + page_table, + request_indices, + context_lens, + output, + validation_scope=object(), + ) + assert scheduler_call_count == 2 * int(scheduler_op is not None) + expected_rows = [] + for batch_index in range(batch_size): + length = int(context_lens[batch_index].item()) + if length == 0: + expected_rows.append(torch.zeros_like(q_latent[batch_index])) + continue + row = int(request_indices[batch_index].item()) + active = page_table[row, :length].long() + logits = q_rope[batch_index].float() @ rope_cache[active, 0].float().T + logits += q_latent[batch_index].float() @ latent_cache[active, 0].float().T + probs = torch.softmax(logits * (256**-0.5), dim=-1) + expected_rows.append( + (probs @ latent_cache[active, 0].float()).to(torch.bfloat16) + ) + expected = torch.stack(expected_rows) + + torch.testing.assert_close(actual, expected, rtol=3e-2, atol=3e-2) + graph = torch.cuda.CUDAGraph() + graph_output = torch.empty_like(output) + with torch.cuda.graph(graph): + kernel( + q_rope, + q_latent, + rope_cache, + latent_cache, + page_table, + request_indices, + context_lens, + graph_output, + validation_scope=object(), + ) + second_graph = torch.cuda.CUDAGraph() + second_graph_output = torch.empty_like(output) + with torch.cuda.graph(second_graph): + kernel( + q_rope, + q_latent, + rope_cache, + latent_cache, + page_table, + request_indices, + context_lens, + second_graph_output, + validation_scope=object(), + ) + graph.replay() + second_graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(graph_output, expected, rtol=3e-2, atol=3e-2) + torch.testing.assert_close( + second_graph_output, + expected, + rtol=3e-2, + atol=3e-2, + ) + + +@pytest.mark.skipif( + not torch.cuda.is_available() or not sgl_fa3_support()[0], + reason="CUDA and a validated sgl-kernel are required", +) +def test_sgl_fa3_varlen_latent_prefill_matches_causal_torch() -> None: + torch.manual_seed(20260807) + device = torch.device("cuda") + heads, width = 10, 8 + chunk_lens = (3, 2) + context_lens = torch.tensor([5, 4], device=device, dtype=torch.int32) + cu_seqlens_q = torch.tensor([0, 3, 5], device=device, dtype=torch.int32) + query_tokens = int(cu_seqlens_q[-1].item()) + q_rope = torch.randn( + query_tokens, heads, 64, device=device, dtype=torch.bfloat16 + ) + q_latent = torch.randn( + query_tokens, heads, 512, device=device, dtype=torch.bfloat16 + ) + rope_cache = torch.randn( + 2 * width, 1, 64, device=device, dtype=torch.bfloat16 + ) + latent_cache = torch.randn( + 2 * width, 1, 512, device=device, dtype=torch.bfloat16 + ) + page_table = torch.arange( + 2 * width, device=device, dtype=torch.int32 + ).view(2, width) + request_indices = torch.tensor([1, 0], device=device, dtype=torch.int32) + output = torch.empty_like(q_latent) + kernel = SglFa3DecodeKernel( + device=device, + max_batch_size=2, + softmax_scale=256**-0.5, + ) + + kernel.run_varlen( + q_rope, + q_latent, + rope_cache, + latent_cache, + page_table, + request_indices, + context_lens, + output, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=max(chunk_lens), + ) + expected_rows = [] + query_start = 0 + for batch_index, chunk_len in enumerate(chunk_lens): + context_len = int(context_lens[batch_index].item()) + row = int(request_indices[batch_index].item()) + for query_offset in range(chunk_len): + visible_len = context_len - chunk_len + query_offset + 1 + active = page_table[row, :visible_len].long() + query_index = query_start + query_offset + logits = q_rope[query_index].float() @ rope_cache[active, 0].float().T + logits += q_latent[query_index].float() @ latent_cache[active, 0].float().T + probs = torch.softmax(logits * (256**-0.5), dim=-1) + expected_rows.append( + (probs @ latent_cache[active, 0].float()).to(torch.bfloat16) + ) + query_start += chunk_len + + torch.testing.assert_close( + output, + torch.stack(expected_rows), + rtol=3e-2, + atol=3e-2, + ) + + +@pytest.mark.skipif( + not torch.cuda.is_available() or not sgl_fa3_support()[0], + reason="CUDA and a validated sgl-kernel are required", +) +def test_sgl_fa3_varlen_explicit_prefill_matches_causal_torch() -> None: + torch.manual_seed(20260807) + device = torch.device("cuda") + heads, width, head_dim = 10, 8, 256 + chunk_lens = (3, 2) + context_lens = torch.tensor([5, 4], device=device, dtype=torch.int32) + cu_seqlens_q = torch.tensor([0, 3, 5], device=device, dtype=torch.int32) + query_tokens = int(cu_seqlens_q[-1].item()) + q = torch.randn( + query_tokens, heads, head_dim, device=device, dtype=torch.bfloat16 + ) + k_cache = torch.randn( + 2 * width, heads, head_dim, device=device, dtype=torch.bfloat16 + ) + v_cache = torch.randn_like(k_cache) + page_table = torch.arange( + 2 * width, device=device, dtype=torch.int32 + ).view(2, width) + request_indices = torch.tensor([1, 0], device=device, dtype=torch.int32) + output = torch.empty_like(q) + kernel = SglFa3DecodeKernel( + device=device, + max_batch_size=2, + softmax_scale=head_dim**-0.5, + ) + + kernel.run_explicit_varlen( + q, + k_cache, + v_cache, + page_table, + request_indices, + context_lens, + output, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=max(chunk_lens), + ) + expected_rows = [] + query_start = 0 + for batch_index, chunk_len in enumerate(chunk_lens): + context_len = int(context_lens[batch_index].item()) + row = int(request_indices[batch_index].item()) + for query_offset in range(chunk_len): + visible_len = context_len - chunk_len + query_offset + 1 + active = page_table[row, :visible_len].long() + query_index = query_start + query_offset + logits = torch.einsum( + "hd,lhd->hl", + q[query_index].float(), + k_cache[active].float(), + ) + probs = torch.softmax(logits * (head_dim**-0.5), dim=-1) + expected_rows.append( + torch.einsum("hl,lhd->hd", probs, v_cache[active].float()).to( + torch.bfloat16 + ) + ) + query_start += chunk_len + + torch.testing.assert_close( + output, + torch.stack(expected_rows), + rtol=3e-2, + atol=3e-2, + ) diff --git a/tests/test_sgl_moe.py b/tests/test_sgl_moe.py new file mode 100644 index 00000000..25624573 --- /dev/null +++ b/tests/test_sgl_moe.py @@ -0,0 +1,154 @@ +from __future__ import annotations + +import importlib.metadata +from unittest.mock import patch + +import pytest +import torch + +from sparsevllm.operators.sgl_moe import ( + SglGlmFusedMoeGate, + _parse_version, + sgl_fused_moe_gate_support, + sgl_moe_alignment_support, + sgl_moe_align_block_size, +) +from sparsevllm.triton_kernel.moe import moe_align_block_size + + +def test_sgl_moe_version_parser_accepts_post_release() -> None: + assert _parse_version("0.3.14.post1+cu128") == (0, 3, 14) + + +def test_sgl_moe_support_rejects_missing_package() -> None: + with patch("importlib.util.find_spec", return_value=None): + assert sgl_fused_moe_gate_support() == ( + False, + "sgl-kernel is not installed", + ) + + +def test_sgl_moe_support_accepts_sglang_kernel_04_package() -> None: + def package_version(distribution: str) -> str: + if distribution == "sglang-kernel": + return "0.4.5" + raise importlib.metadata.PackageNotFoundError(distribution) + + with ( + patch("importlib.util.find_spec", return_value=object()), + patch("importlib.metadata.version", side_effect=package_version), + ): + gate_supported, gate_reason = sgl_fused_moe_gate_support() + align_supported, align_reason = sgl_moe_alignment_support() + + assert gate_supported + assert align_supported + assert "0.4.5" in gate_reason + assert "0.4.5" in align_reason + + +@pytest.mark.skipif( + not torch.cuda.is_available() or not sgl_fused_moe_gate_support()[0], + reason="CUDA and a validated sgl-kernel are required", +) +def test_sgl_glm_fused_gate_matches_reference_and_replays_graph() -> None: + torch.manual_seed(20260807) + logits = torch.randn(32, 64, dtype=torch.float32, device="cuda") * 3 + correction_bias = ( + torch.randn(64, dtype=torch.float32, device="cuda") * 0.1 + ) + scaling = 1.8 + scores = torch.sigmoid(logits) + reference_ids = torch.topk( + scores + correction_bias, + 4, + dim=-1, + sorted=False, + ).indices + reference_weights = scores.gather(1, reference_ids) + reference_weights /= reference_weights.sum(dim=-1, keepdim=True) + 1e-20 + reference_weights *= scaling + gate = SglGlmFusedMoeGate(num_experts=64, top_k=4) + + actual_weights, actual_ids = gate( + logits, + correction_bias, + top_k=4, + routed_scaling_factor=scaling, + ) + + reference_order = reference_ids.argsort(dim=-1) + actual_order = actual_ids.argsort(dim=-1) + sorted_reference_ids = reference_ids.gather(1, reference_order) + sorted_actual_ids = actual_ids.gather(1, actual_order) + sorted_reference_weights = reference_weights.gather(1, reference_order) + sorted_actual_weights = actual_weights.gather(1, actual_order) + assert torch.equal(sorted_actual_ids, sorted_reference_ids) + torch.testing.assert_close( + sorted_actual_weights, + sorted_reference_weights, + rtol=2e-6, + atol=2e-7, + ) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_weights, graph_ids = gate( + logits, + correction_bias, + top_k=4, + routed_scaling_factor=scaling, + ) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(graph_ids, actual_ids) + torch.testing.assert_close(graph_weights, actual_weights, rtol=0, atol=0) + + +@pytest.mark.skipif( + not torch.cuda.is_available() or not sgl_fused_moe_gate_support()[0], + reason="CUDA and a validated sgl-kernel are required", +) +def test_sgl_moe_alignment_matches_grouped_assignments() -> None: + torch.manual_seed(20260808) + topk_ids = torch.randint( + 0, + 64, + (32, 4), + dtype=torch.int32, + device="cuda", + ) + topk_ids[0].fill_(63) + reference_sorted, reference_experts, reference_count = moe_align_block_size( + topk_ids, + 16, + 64, + ) + actual = sgl_moe_align_block_size( + topk_ids, + block_size=16, + num_experts=64, + local_expert_start=0, + local_expert_end=64, + ) + torch.cuda.synchronize() + count = int(reference_count.item()) + assert int(actual.num_tokens_post_padded.item()) == count + assert torch.equal( + actual.expert_ids[: count // 16], + reference_experts[: count // 16], + ) + num_assignments = int(topk_ids.numel()) + for block_index in range(count // 16): + block = slice(block_index * 16, (block_index + 1) * 16) + reference_ids = sorted( + value + for value in reference_sorted[block].tolist() + if value < num_assignments + ) + actual_ids = sorted( + value + for value in actual.sorted_token_ids[block].tolist() + if value < num_assignments + ) + assert actual_ids == reference_ids diff --git a/tests/test_tp_rpc.py b/tests/test_tp_rpc.py index c5a05448..31343c53 100644 --- a/tests/test_tp_rpc.py +++ b/tests/test_tp_rpc.py @@ -450,7 +450,8 @@ def test_model_runner_exit_drains_graphs_before_barrier(): unlink=lambda: calls.append("unlink_shm"), ) runner.parallel_context = SimpleNamespace( - world_barrier=lambda **_: calls.append("barrier") + world_barrier=lambda **_: calls.append("barrier"), + close_all_reduce_provider=lambda: calls.append("close_all_reduce"), ) with ( @@ -472,6 +473,7 @@ def test_model_runner_exit_drains_graphs_before_barrier(): "close_shm", "barrier", "unlink_shm", + "close_all_reduce", "reset", "destroy", ] diff --git a/tests/test_triton_moe.py b/tests/test_triton_moe.py index a0c29f5a..fdd6cd7f 100644 --- a/tests/test_triton_moe.py +++ b/tests/test_triton_moe.py @@ -8,11 +8,18 @@ from sparsevllm.triton_kernel.gate_up_swiglu import h20_gate_up_swiglu from sparsevllm.triton_kernel.moe import ( _prepare_expert_assignment, + append_shared_expert_route, fused_moe, fused_moe_gate_up_swiglu, moe_align_block_size, ) from sparsevllm.triton_kernel.moe_topk import topk_softmax +from sparsevllm.triton_kernel.silu_and_mul import _resolve_silu_launch_config + + +def test_silu_launch_config_uses_decode_tile_only_for_small_rows(): + assert _resolve_silu_launch_config(160) == (32, 128, 4) + assert _resolve_silu_launch_config(257) == (128, 128, None) def _is_h20() -> bool: @@ -89,6 +96,48 @@ def _oracle_local_moe( return output +@unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for Triton MoE tests.") +def test_append_shared_expert_route_matches_cat_and_replays_graph(): + ids = torch.tensor( + [[3, 1, 7, 2], [9, 8, 4, 6]], + dtype=torch.int32, + device="cuda", + ) + weights = torch.tensor( + [[0.1, 0.2, 0.3, 0.4], [0.4, 0.3, 0.2, 0.1]], + dtype=torch.float32, + device="cuda", + ) + actual_ids, actual_weights = append_shared_expert_route( + ids, + weights, + shared_expert_id=64, + ) + expected_ids = torch.cat( + (ids, torch.full_like(ids[:, :1], 64)), + dim=1, + ) + expected_weights = torch.cat( + (weights, torch.ones_like(weights[:, :1])), + dim=1, + ) + torch.cuda.synchronize() + assert torch.equal(actual_ids, expected_ids) + assert torch.equal(actual_weights, expected_weights) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_ids, graph_weights = append_shared_expert_route( + ids, + weights, + shared_expert_id=64, + ) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(graph_ids, expected_ids) + assert torch.equal(graph_weights, expected_weights) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for Triton MoE tests.") def test_moe_align_block_size_filters_ep_experts_and_pads_blocks(): topk_ids = torch.tensor( diff --git a/tests/test_validate_glm_checkpoint_parallel.py b/tests/test_validate_glm_checkpoint_parallel.py index 7e6ee4b8..6bc58bfe 100644 --- a/tests/test_validate_glm_checkpoint_parallel.py +++ b/tests/test_validate_glm_checkpoint_parallel.py @@ -1,12 +1,46 @@ from __future__ import annotations +import subprocess +from pathlib import Path + import pytest from scripts.validation.validate_glm_checkpoint_parallel import ( + REPO_ROOT, + _git_output, _validate_parallel_evidence, ) +def test_git_provenance_is_resolved_from_repository_root( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + monkeypatch.chdir(tmp_path) + + assert _git_output("rev-parse", "--show-toplevel") == str(REPO_ROOT) + + +def test_git_provenance_failure_is_explicit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def fail(*args, **kwargs): + del args, kwargs + raise subprocess.CalledProcessError( + returncode=128, + cmd=["git", "rev-parse", "HEAD"], + stderr="fatal: unavailable repository", + ) + + monkeypatch.setattr(subprocess, "run", fail) + + with pytest.raises( + RuntimeError, + match="Git provenance command failed.*fatal: unavailable repository", + ): + _git_output("rev-parse", "HEAD") + + def _graph_summary(enabled: bool = False) -> dict: return { "enabled": enabled, From cae8e2b3e2431ea4f8b5c5330baf50ac2cef6dca Mon Sep 17 00:00:00 2001 From: CURRENTF <1206013760@qq.com> Date: Sun, 9 Aug 2026 04:09:38 +0800 Subject: [PATCH 18/48] perf: checkpoint glm flash runtime --- .../optimize-sparsevllm-kernel/SKILL.md | 160 ++++++ .../agents/openai.yaml | 4 + .../references/benchmark-protocol.md | 74 +++ .../references/cuda-cute.md | 39 ++ .../references/nsight-playbook.md | 62 +++ .../references/operator-integration.md | 53 ++ .../references/reference-sources.md | 36 ++ .../references/tilelang.md | 52 ++ .../references/triton.md | 53 ++ AGENTS.md | 5 +- README.md | 54 +- docs/en/getting_started/README.md | 48 +- docs/en/getting_started/reproducibility.md | 20 +- docs/zh/getting_started/README.md | 47 +- docs/zh/getting_started/reproducibility.md | 16 +- pyproject.toml | 43 +- requirements/locks/README.md | 29 + requirements/locks/canonical-cu129-py310.txt | 149 +++++ src/sparsevllm/operators/mla_attention.py | 149 +++++ src/sparsevllm/operators/tilelang_mla.py | 516 ++++++++++++++++++ src/sparsevllm/tilelang_kernel/__init__.py | 1 + .../tilelang_kernel/mla/LICENSE.tilelang | 21 + src/sparsevllm/tilelang_kernel/mla/README.md | 28 + .../tilelang_kernel/mla/__init__.py | 5 + src/sparsevllm/tilelang_kernel/mla/decode.py | 406 ++++++++++++++ .../utils/select_omnikv_full_layers.py | 92 +++- tests/test_dependency_constraints.py | 62 ++- tests/test_mla_attention_operator.py | 33 +- tests/test_omnikv_full_layer_selector.py | 38 ++ tests/test_tilelang_mla_kernel.py | 205 +++++++ tests/test_tilelang_mla_operator.py | 439 +++++++++++++++ 31 files changed, 2768 insertions(+), 171 deletions(-) create mode 100644 .agents/skills/optimize-sparsevllm-kernel/SKILL.md create mode 100644 .agents/skills/optimize-sparsevllm-kernel/agents/openai.yaml create mode 100644 .agents/skills/optimize-sparsevllm-kernel/references/benchmark-protocol.md create mode 100644 .agents/skills/optimize-sparsevllm-kernel/references/cuda-cute.md create mode 100644 .agents/skills/optimize-sparsevllm-kernel/references/nsight-playbook.md create mode 100644 .agents/skills/optimize-sparsevllm-kernel/references/operator-integration.md create mode 100644 .agents/skills/optimize-sparsevllm-kernel/references/reference-sources.md create mode 100644 .agents/skills/optimize-sparsevllm-kernel/references/tilelang.md create mode 100644 .agents/skills/optimize-sparsevllm-kernel/references/triton.md create mode 100644 requirements/locks/README.md create mode 100644 requirements/locks/canonical-cu129-py310.txt create mode 100644 src/sparsevllm/operators/tilelang_mla.py create mode 100644 src/sparsevllm/tilelang_kernel/__init__.py create mode 100644 src/sparsevllm/tilelang_kernel/mla/LICENSE.tilelang create mode 100644 src/sparsevllm/tilelang_kernel/mla/README.md create mode 100644 src/sparsevllm/tilelang_kernel/mla/__init__.py create mode 100644 src/sparsevllm/tilelang_kernel/mla/decode.py create mode 100644 tests/test_tilelang_mla_kernel.py create mode 100644 tests/test_tilelang_mla_operator.py diff --git a/.agents/skills/optimize-sparsevllm-kernel/SKILL.md b/.agents/skills/optimize-sparsevllm-kernel/SKILL.md new file mode 100644 index 00000000..1bf4c01c --- /dev/null +++ b/.agents/skills/optimize-sparsevllm-kernel/SKILL.md @@ -0,0 +1,160 @@ +--- +name: optimize-sparsevllm-kernel +description: Optimize and integrate Sparse-vLLM GPU kernels across Triton, TileLang, CUDA/CuTe, and external SGL kernels. Use when Codex needs to identify an LLM inference hotspot, choose a kernel implementation path, write or fuse a kernel, tune an existing kernel, build correctness or microbenchmark coverage, analyze Nsight Compute results, integrate a provider, or validate kernel and end-to-end performance. +--- + +# Optimize Sparse-vLLM Kernel + +Use one evidence loop from the serving workload to the kernel and back: + +```text +matched LLM workload + -> hotspot evidence + -> implementation choice + -> correctness oracle + -> stable microbenchmark + -> targeted tuning + -> Nsight Compute + -> operator/provider integration + -> matched end-to-end validation +``` + +Do not assume that a kernel rewrite is useful before locating its contribution +to the requested workload. If the user already specifies a kernel, proceed but +state whether end-to-end hotspot evidence exists. + +## Load the Relevant Guidance + +Read [benchmark-protocol.md](references/benchmark-protocol.md) before running +any performance experiment. + +Read exactly one implementation guide first, then load another only when a +measured comparison requires it: + +- Triton: [triton.md](references/triton.md) +- TileLang: [tilelang.md](references/tilelang.md) +- CUDA, CuTe, CUTLASS, or an external compiled kernel: + [cuda-cute.md](references/cuda-cute.md) + +Read [nsight-playbook.md](references/nsight-playbook.md) before collecting or +interpreting Nsight Compute data. Read +[operator-integration.md](references/operator-integration.md) before changing +provider selection, dependencies, workspaces, layouts, model call sites, or +production dispatch. + +Use [reference-sources.md](references/reference-sources.md) only when an +upstream example is needed. Pin the exact source revision and inspect its +license before adapting code. + +When available, use companion skills for their focused expertise: + +- `llm-torch-profiler-analysis` for SGLang/vLLM/TRT-LLM trace analysis. +- `kernel-triton-writing` for Triton implementation details. +- `add-jit-kernel` for JIT CUDA integration without a large C++ project. +- `add-sgl-kernel` for CUTLASS or complex AOT integration. +- `kernel-cute-writing` for CuTe-specific implementation. +- `perf-nsight-compute-analysis` for deep Nsight Compute interpretation. +- `debug-cuda-crash` for illegal access, misalignment, or graph-capture faults. + +Do not block when a companion skill is unavailable; follow this skill's local +references and report the missing capability. + +## Execute the Workflow + +### 1. Freeze the Scope + +Inspect Git status before editing and preserve unrelated tracked and untracked +work. Define the operator, phase, workload, hardware, dtype, shapes, layouts, +parallel topology, CUDA Graph mode, and comparison baseline. Check all devices +and select an idle permitted GPU before starting a GPU task. + +### 2. Locate the Cost + +Profile a representative end-to-end workload. Separate prefill, decode, +sampling, communication, host overhead, graph replay, and compilation. Rank +hotspots by total contribution rather than kernel latency alone. Record fusion +and overlap opportunities, but treat them as hypotheses until measured. + +### 3. Choose the Implementation Path + +Prefer the smallest path that can express the required computation: + +- Keep or improve Triton for broadly applicable repository-owned kernels. +- Use TileLang when explicit tiling, pipelining, shared-memory layouts, + tensor-core scheduling, TMA, or warp specialization materially helps. +- Prefer JIT CUDA when custom CUDA is needed without CUTLASS or a large AOT + project. +- Use compiled SGL/CUTLASS/CuTe integration only when the required primitives, + layouts, or performance cannot be reached cleanly with a JIT path. + +Do not replace a mature provider solely because another DSL looks promising. + +### 4. Establish Correctness + +Create an independent Torch or mathematically direct oracle. Specify input and +output shapes, dtypes, strides, aliases, mutation, padding, empty cases, +numerical tolerances, and reduction semantics. Cover real model shapes, +boundary shapes, non-contiguous inputs when supported, repeated execution, and +CUDA Graph capture/replay when claimed. + +Reject a candidate immediately when correctness fails. Never tune against a +known-wrong implementation. + +### 5. Build the Microbenchmark + +Benchmark the actual callable boundary used by serving, including required +workspace initialization, output reset, synchronization, or materialization. +Separate compile and cold-start cost from steady-state latency. Use shapes +derived from the workload, fixed inputs and seeds, sufficient warmup, raw +samples, and identical conditions for baseline and candidate. + +### 6. Tune with Bounded Hypotheses + +Change one explained dimension or run a declared finite matrix. Record every +candidate, including failures. Typical dimensions include tile shape, program +grid, threads or warps, pipeline stages, split count, vector width, +shared-memory layout, swizzle, async copy/TMA, fusion boundary, register use, +and workspace layout. + +Tune offline. Do not benchmark or search configurations in the serving hot +path. Bind the selected configuration deterministically from validated shape +and device facts. + +### 7. Profile Finalists + +Run Nsight Compute only after correctness and stable timing identify a small +set of finalists. Compare the same shape and call boundary. Use measured +occupancy, achieved bandwidth, tensor-core utilization, register or local +memory use, shared-memory behavior, scheduler issue rate, and warp stalls to +support the next hypothesis. Do not infer a bottleneck from occupancy alone. + +### 8. Integrate through Operators + +Keep kernels under the repository-owned kernel packages and keep dependency +checks, support predicates, workspace ownership, launch selection, and +fallback policy under `src/sparsevllm/operators/`. Resolve and bind a provider +before the forward hot path. Route unsupported configurations before launch; +never catch a runtime kernel failure and silently switch providers. + +### 9. Return to End-to-End Measurement + +Repeat the original workload with the same model, request trace, concurrency, +context/output lengths, TP/EP topology, cache state, graph mode, and metric +window. Report kernel latency improvement separately from end-to-end latency, +throughput, TTFT, or TPOT. A faster microbenchmark is not an end-to-end win. + +## Acceptance Gates + +Require all applicable gates before calling the work complete: + +1. Independent correctness equivalence passes. +2. Provider selection and rejection paths are tested. +3. Claimed devices, dtypes, shapes, and graph modes are exercised on hardware. +4. Microbenchmark raw samples and summary are saved. +5. Nsight evidence exists for hardware-level bottleneck claims. +6. Matched end-to-end validation supports serving-level claims. +7. Commands, Git state, environment, selected provider, and artifacts are + recorded. + +Mark any unrun gate explicitly. Preserve the best verified implementation and +the baseline; do not leave an unverified candidate as the production default. diff --git a/.agents/skills/optimize-sparsevllm-kernel/agents/openai.yaml b/.agents/skills/optimize-sparsevllm-kernel/agents/openai.yaml new file mode 100644 index 00000000..65bf1692 --- /dev/null +++ b/.agents/skills/optimize-sparsevllm-kernel/agents/openai.yaml @@ -0,0 +1,4 @@ +interface: + display_name: "Optimize Sparse-vLLM Kernel" + short_description: "Optimize Sparse-vLLM GPU kernels end to end" + default_prompt: "Use $optimize-sparsevllm-kernel to find, implement, tune, profile, and integrate a Sparse-vLLM GPU kernel." diff --git a/.agents/skills/optimize-sparsevllm-kernel/references/benchmark-protocol.md b/.agents/skills/optimize-sparsevllm-kernel/references/benchmark-protocol.md new file mode 100644 index 00000000..c7517659 --- /dev/null +++ b/.agents/skills/optimize-sparsevllm-kernel/references/benchmark-protocol.md @@ -0,0 +1,74 @@ +# Kernel Benchmark Protocol + +Apply this protocol to every performance claim. + +## Before Launch + +1. Inspect every GPU's utilization, memory, and compute-process ownership. +2. Select an idle permitted device. If none is idle, wait or report instead of + sharing a busy device. +3. Inspect Git status and preserve unrelated tracked and untracked work. +4. Record the kernel callable boundary and the end-to-end workload that + produced the target shapes. +5. Establish correctness before collecting performance samples. + +## Record the Case + +Save a manifest containing: + +- repository path, Git SHA, branch, and dirty status +- exact command and interpreter +- Torch, Triton, TileLang, CUDA, driver, and relevant kernel-package versions +- GPU name, compute capability, selected device, clocks or power constraints + when controlled +- input shapes, dtypes, strides, seed, and data-generation method +- provider, launch configuration, graph mode, TP/EP topology, and cache state +- warmup count, timed repetitions, timing method, and synchronization boundary +- output paths and explicit run status + +Do not overwrite previous raw results. Use a new run directory or immutable +case identifier. + +## Microbenchmark Semantics + +- Time the serving-relevant wrapper, not only an internal launch, unless the + result is explicitly labeled kernel-only. +- Include required output reset, workspace preparation, conversion, or + materialization. Report optional components separately when decomposition is + useful. +- Exclude compilation from steady-state latency after recording cold compile + time separately. +- Warm every specialization and synchronize before and after the timed region. +- Use CUDA events or another GPU-aware timer correctly; never time asynchronous + launches with host wall time alone. +- Use identical inputs, streams, graph mode, and synchronization for baseline + and candidate. Interleave them when long runs may drift. +- Save raw samples and report at least sample count and median. Add tail or + dispersion statistics when they affect the decision. +- State whether caches are intentionally warm or cold. Do not mix regimes. +- Repeat suspicious gains and reject results affected by competing processes, + throttling, compilation, or changing clocks. + +## End-to-End Semantics + +Match model/checkpoint, request trace, prompt and output lengths, batch and +concurrency, TP/EP, cache state, graph/provider settings, decoding parameters, +and metric window. Keep TTFT, TPOT/inter-token latency, request latency, +throughput, and kernel time distinct. Label partial serving runs separately +from completed end-to-end results. + +## Artifact Minimum + +Persist: + +```text +run_manifest.json +raw_samples.jsonl +summary.json +stdout.log +stderr.log +``` + +Add profiler traces and `.ncu-rep` files when collected. Use structured status +and error fields; do not treat a non-empty output or a successful process start +as benchmark success. diff --git a/.agents/skills/optimize-sparsevllm-kernel/references/cuda-cute.md b/.agents/skills/optimize-sparsevllm-kernel/references/cuda-cute.md new file mode 100644 index 00000000..ed3be74f --- /dev/null +++ b/.agents/skills/optimize-sparsevllm-kernel/references/cuda-cute.md @@ -0,0 +1,39 @@ +# CUDA, CuTe, CUTLASS, and SGL Kernel Guide + +Use a compiled or JIT CUDA path only when its required primitives, layout +control, or measured performance justify the added integration cost. + +## Choose the Smallest Integration + +- Prefer JIT CUDA when the kernel does not require CUTLASS or a large C++ + project. Use `add-jit-kernel` when available. +- Use CuTe when explicit NVIDIA tensor layouts and architecture-specific + primitives are central. Use `kernel-cute-writing` when available. +- Use SGL/CUTLASS AOT integration for complex compiled projects or packaged + kernels. Use `add-sgl-kernel` when available. +- Keep a verified Triton or Torch implementation as the correctness baseline + and, where supported, the portable provider. + +## Control the Contract + +Declare supported compute capabilities, CUDA/toolchain versions, dtypes, +alignments, layouts, workspace, streams, graph behavior, and mutation. Keep +build and package availability checks inside the provider and import compiled +extensions lazily. + +Do not expose CUTLASS packing, reordered projections, descriptor formats, or +workspace details to model classes. The selected provider owns physical +layouts and preparation. Reject unsupported configurations during resolution +or preparation; do not catch a launch failure and switch implementations. + +## Validate + +Compare the actual wrapper against an independent oracle over production and +boundary shapes. Exercise the minimum declared dependency/toolchain and every +claimed architecture on real hardware. Check sanitizer or crash diagnostics +for indexing and lifetime changes, and validate CUDA Graph capture/replay when +advertised. + +Measure compile/startup separately from steady-state execution. Include any +descriptor construction, packing, workspace clearing, or output conversion +that remains in the service call boundary. diff --git a/.agents/skills/optimize-sparsevllm-kernel/references/nsight-playbook.md b/.agents/skills/optimize-sparsevllm-kernel/references/nsight-playbook.md new file mode 100644 index 00000000..57df19be --- /dev/null +++ b/.agents/skills/optimize-sparsevllm-kernel/references/nsight-playbook.md @@ -0,0 +1,62 @@ +# Nsight Compute Playbook + +Use Nsight Compute after correctness and stable microbenchmarks narrow the +candidate set. Nsight replay and metric collection can perturb execution, so +do not use profiler duration as the production latency result. + +## Capture Deliberately + +1. Choose one representative shape and a callable that launches the target + kernel predictably. +2. Warm JIT compilation before capture. +3. Filter the target kernel or limit launches when possible. +4. Start with focused sections, then collect a broader set only when needed. +5. Save the command, console output, and `.ncu-rep` artifact. +6. Capture baseline and candidate under the same software and hardware state. + +Useful section families include Speed of Light, Launch Statistics, Occupancy, +Memory Workload Analysis, Scheduler Statistics, Warp State Statistics, and +Source Counters. Confirm the exact section names supported by the installed +`ncu` version before scripting them. + +## Interpret as a Chain of Evidence + +### Launch and Occupancy + +Check grid size, waves per SM, block/thread shape, shared memory, registers per +thread, theoretical occupancy, and achieved occupancy. Low occupancy matters +only when it limits latency hiding or parallelism; high occupancy does not +prove efficiency. + +### Memory + +Compare achieved DRAM/L2/shared throughput, transaction efficiency, cache hit +rates, sectors, and shared-bank conflicts. Relate bytes moved to the algorithm +and wrapper, including intermediate tensors eliminated or introduced by +fusion. + +### Compute + +Check tensor-core or arithmetic-pipe utilization, instruction mix, issue rate, +and dependency stalls. Verify that the chosen tile and dtype actually reach +the intended hardware path. + +### Registers and Stalls + +Inspect register count, local-memory traffic, spills, scoreboard/dependency +stalls, barrier stalls, memory throttling, and not-selected warps. Connect a +stall change to a concrete code or schedule change before acting on it. + +### Roofline + +Estimate arithmetic intensity using the measured callable boundary. Classify +memory- versus compute-limited behavior only when achieved bandwidth/compute +and the traffic model agree. Re-evaluate after fusion because the boundary and +bytes moved have changed. + +## Close the Loop + +Use the profile to form one next hypothesis, benchmark the resulting change, +and retain it only if stable latency improves without breaking correctness. +Do not optimize a metric that does not move the microbenchmark or end-to-end +result. diff --git a/.agents/skills/optimize-sparsevllm-kernel/references/operator-integration.md b/.agents/skills/optimize-sparsevllm-kernel/references/operator-integration.md new file mode 100644 index 00000000..fc997272 --- /dev/null +++ b/.agents/skills/optimize-sparsevllm-kernel/references/operator-integration.md @@ -0,0 +1,53 @@ +# Sparse-vLLM Operator Integration + +Read the sibling +[`review-operator-organization`](../../review-operator-organization/SKILL.md) +skill before changing production dispatch. Treat its architecture and +validation rules as authoritative. + +## Keep Ownership Explicit + +```text +model semantic call + -> OpSpec + -> OpResolver(DeviceCaps) + -> selected OperatorProvider + -> provider-owned preparation/workspace + -> Triton, TileLang, JIT, or external kernel +``` + +- Keep repository-owned Triton kernels under `src/sparsevllm/triton_kernel/`. +- Keep repository-owned TileLang kernels under + `src/sparsevllm/tilelang_kernel/`. +- Keep support checks, dependency availability, static launch selection, + workspaces, and provider binding under `src/sparsevllm/operators/`. +- Keep models expressed in semantic operations rather than backend names, + package imports, device probes, or physical weight layouts. +- Keep device discovery and stable capability facts under the platform layer. + +## Resolve Before Execution + +Make `supports(spec, caps)` cover platform, architecture, dtype, quantization, +shape and alignment, layouts, topology, graph behavior, workspace, toolchain, +and external API availability. Import optional compilers and kernel packages +lazily. Bind the provider outside the forward hot path. + +Allow fallback only by rejecting an unsupported provider during resolution or +preparation. Once execution begins, surface compilation and launch failures. +Do not make one provider failure disable unrelated operators. + +## Validate the Complete Path + +Add or update: + +- independent kernel-equivalence tests +- resolver selection and rejection tests +- missing and minimum-version dependency tests +- boundary shape, dtype, stride, padding, and workspace tests +- CUDA Graph capture/replay tests where supported +- actual model-path integration coverage +- matched performance evidence for every priority change + +Record the selected provider and rejection reasons in reproducible artifacts. +After implementation, use `$review-operator-organization` for the focused +architecture review and `$code-review` for the complete diff. diff --git a/.agents/skills/optimize-sparsevllm-kernel/references/reference-sources.md b/.agents/skills/optimize-sparsevllm-kernel/references/reference-sources.md new file mode 100644 index 00000000..4a91c08a --- /dev/null +++ b/.agents/skills/optimize-sparsevllm-kernel/references/reference-sources.md @@ -0,0 +1,36 @@ +# Kernel Reference Sources + +Use upstream code as design evidence and few-shot material, not as an implicit +dependency or an unquestioned performance baseline. Before adapting code, +record the exact commit, source path, license, local changes, supported +hardware, and benchmark protocol. + +## TileLang + +- [Tile-AI/TileLang](https://github.com/tile-ai/tilelang): inspect official + examples and documentation for the installed version. The DeepSeek MLA + examples are especially relevant to tiled attention, split-KV, layouts, + pipelining, shared-memory swizzle, and warp specialization. +- [DeepSeek TileKernels](https://github.com/deepseek-ai/TileKernels): use as a + kernel portfolio and testing reference. Treat each kernel's hardware, + dependency, and benchmark assumptions as local to its pinned revision. + +## Serving Integration + +- [SGLang](https://github.com/sgl-project/sglang): inspect current in-tree + Triton and TileLang serving kernels, provider/backend dispatch, graph + constraints, tests, and benchmarks. Search the pinned checkout rather than + relying on remembered paths because the tree changes frequently. +- Inspect SGLang's JIT kernel path before introducing a standalone C++ project. + Inspect the separate SGL kernel package only when CUTLASS, CuTe, AOT build, + or packaged binary integration is required. + +## Reuse Rules + +1. Prefer a repository-owned or installed-version example over `main`. +2. Compare tensor semantics, layouts, dtypes, scaling, masking, and mutation + before comparing implementation shape. +3. Port the smallest relevant mechanism rather than copying an entire module. +4. Retain required license and provenance files and describe local changes. +5. Rebuild correctness and benchmark baselines under Sparse-vLLM's actual + serving contract. diff --git a/.agents/skills/optimize-sparsevllm-kernel/references/tilelang.md b/.agents/skills/optimize-sparsevllm-kernel/references/tilelang.md new file mode 100644 index 00000000..7db96924 --- /dev/null +++ b/.agents/skills/optimize-sparsevllm-kernel/references/tilelang.md @@ -0,0 +1,52 @@ +# TileLang Kernel Guide + +Use this guide for kernels under `src/sparsevllm/tilelang_kernel/` and their +lazy adapters under `src/sparsevllm/operators/`. + +## Establish the Contract + +Start from a direct Torch oracle and, when useful, a validated Triton provider. +Record the exact TileLang, TVM-FFI, Torch, CUDA, driver, and GPU architecture +versions. Validate the installed versions used by the production dependency +contract; do not rely only on a newer developer checkout. + +Keep source files stable and present while compiling. TileLang JIT may inspect +Python source, so do not benchmark code stored only in a transient shell body +or a file that another worktree operation can remove. Warm every shape needed +by CUDA Graph before capture. + +## Tune Systematically + +Classify the candidate as launch-, memory-, latency-, or compute-sensitive, +then explore a bounded subset of: + +1. Tile shapes along token, head, and hidden dimensions. +2. Block, thread, warp, or warpgroup mapping. +3. Pipeline stage count and the shared-memory cost of each stage. +4. Shared-memory layout, padding, bank conflicts, and swizzle. +5. Async copy or TMA eligibility, alignment, and producer/consumer overlap. +6. Tensor-core tile compatibility and achieved utilization. +7. Fragment size, live ranges, register pressure, spills, and occupancy. +8. Split-KV or split-reduction parallelism and combine-kernel overhead. +9. Fusion benefit versus added score reset, workspace, or atomic traffic. +10. Launch shape and CUDA Graph replay behavior. + +Measure after every change or declared matrix. Keep raw results for losing +configurations so the same search is not repeated. Preserve the best verified +version rather than the last attempted version. + +## Integrate Safely + +Keep the TileLang module limited to kernel definitions. Put device and package +checks, shape support, lazy compilation, output/workspace ownership, static +launch-config selection, and routing under the operator provider. Import +TileLang lazily so an unselected provider does not initialize the compiler. + +Select split counts or other tuned values from offline-calibrated device and +shape buckets. Never benchmark in `forward()`. Route unsupported dtype, layout, +capacity, architecture, or dependency versions to another provider before +launch, and surface kernel compilation or execution failures. + +Validate padded heads, indirect slots, invalid rows, score/output capacity, +atomic reductions, caller-owned workspaces, no-score paths, and CUDA Graph +capture when the kernel supports them. diff --git a/.agents/skills/optimize-sparsevllm-kernel/references/triton.md b/.agents/skills/optimize-sparsevllm-kernel/references/triton.md new file mode 100644 index 00000000..8b948177 --- /dev/null +++ b/.agents/skills/optimize-sparsevllm-kernel/references/triton.md @@ -0,0 +1,53 @@ +# Triton Kernel Guide + +Use this guide for repository-owned Triton kernels and wrappers under +`src/sparsevllm/triton_kernel/`. + +## Inspect Before Editing + +Trace the serving call through its wrapper, launch grid, kernel, output use, +and synchronization boundary. Identify compile-time constants, supported +layouts, mutation, workspace, CUDA Graph constraints, and the actual model +shape distribution. Compare against existing neighboring kernels before +creating a new family. + +Use the external `kernel-triton-writing` skill when it is available. Keep this +guide authoritative for Sparse-vLLM-specific integration and validation. + +## Establish the Baseline + +- Use an independent Torch oracle, not another kernel with the same reduction + or indexing strategy. +- Cover masked tails, empty inputs, odd lengths, large production shapes, + non-contiguous inputs when supported, aliases, and repeated launches. +- Validate the wrapper as well as the JIT function. A correct kernel with an + incorrect grid, stride, or output contract is still wrong. +- Warm every required specialization before CUDA Graph capture and prove replay + stability when graph support is claimed. + +## Tune the Relevant Dimensions + +Measure a bounded matrix chosen from the kernel structure: + +- program decomposition and grid mapping +- block sizes and vector width +- `num_warps` and `num_stages` +- coalescing and redundant loads +- masks and compile-time specialization +- reduction order and accumulation dtype +- fusion versus intermediate tensor traffic +- register pressure, spills, and occupancy +- launch count and graph behavior + +Avoid using private Triton internals for steady-state dispatch. Separate +offline tuning from production launch: tune once for a declared key, then bind +a plain deterministic launcher or static configuration. Never invoke autotune +search inside a request hot path. + +## Interpret Results + +Check whether the wrapper includes extra casts, padding, allocation, output +reset, or synchronization hidden by a kernel-only timer. Preserve numerical +semantics when changing reduction order or precision. Attribute a win to +fusion only when eliminated launches or memory traffic appear in the matched +measurement. diff --git a/AGENTS.md b/AGENTS.md index 09651b0d..f842634a 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -1,12 +1,13 @@ # Repo Skills -This repository includes a repo-local Codex skill. +This repository includes repo-local Codex skills. ## Available skills - `add-sparse-method`: Add or refactor a first-class Sparse-vLLM sparse method following this repo's architecture. Use when Codex needs to introduce a new `vllm_sparse_method`, move method logic out of `attention.py` or `utils/`, add method-specific cache metadata or decode-time view building, and preserve the cache-manager-first design. File: `.agents/skills/add-sparse-method/SKILL.md` - `code-review`: Review Sparse-vLLM diffs for correctness, sparse-runtime and operator architecture, scheduling semantics, reproducibility, performance, and tests. Use when reviewing PRs, git diffs, sparse method integrations, operator/provider or kernel changes, cache-manager or scheduler changes, benchmark/evaluation scripts, OpenAI serving changes, or when the user asks for a code review. File: `.agents/skills/code-review/SKILL.md` - `review-operator-organization`: Review operator/provider boundaries, device capability selection, kernel ownership, dependency compatibility, weight layouts, fallback semantics, and validation. Use for changes under `operators/`, `platforms/`, Triton kernels, external kernel integrations, or model-to-operator call sites. File: `.agents/skills/review-operator-organization/SKILL.md` +- `optimize-sparsevllm-kernel`: Find, implement, tune, profile, and integrate Sparse-vLLM GPU kernels across Triton, TileLang, CUDA/CuTe, and external SGL providers. Use for kernel hotspots, fusion, correctness baselines, microbenchmarks, Nsight Compute analysis, provider integration, or matched end-to-end performance validation. File: `.agents/skills/optimize-sparsevllm-kernel/SKILL.md` ## How to use @@ -14,6 +15,8 @@ This repository includes a repo-local Codex skill. - In this repo, invoke the review skill as `$code-review`. - Invoke focused operator reviews as `$review-operator-organization`; `$code-review` loads it automatically for relevant diffs. +- Invoke the end-to-end kernel workflow as `$optimize-sparsevllm-kernel`; it + loads only the selected DSL and profiling references. - Keep method-specific runtime state in `src/sparsevllm/engine/cache_manager/`. - Keep `src/sparsevllm/layers/attention.py` generic and hook new methods through shared cache-manager interfaces when possible. diff --git a/README.md b/README.md index 74db0d7f..d7505202 100644 --- a/README.md +++ b/README.md @@ -69,7 +69,6 @@ Read the method overview and integration rules in | Qwen3 | ✅ | | Qwen3MoE | ✅ | | Qwen3.5 / Qwen3.6 | ✅ | -| Qwen3.5 / Qwen3.6 MoE | ✅ | | GLM-4.7-Flash | ✅ | | Llama 3 / 3.1 | ✅ | | MiniMax M2.7 | ✅ | @@ -94,8 +93,9 @@ The full documentation index is maintained in [docs/en/README.md](docs/en/README ## Quick Start -Sparse-vLLM requires Python 3.10 or newer. Install the package from the -repository root using the runtime versions pinned in `pyproject.toml`. +Sparse-vLLM requires Python 3.10 or newer. The validated CUDA 12.9 runtime and +test environment is frozen in +`requirements/locks/canonical-cu129-py310.txt`. ### Conda @@ -103,14 +103,9 @@ repository root using the runtime versions pinned in `pyproject.toml`. conda create -n svllm python=3.10 -y conda activate svllm -pip install torch==2.11.0 torchvision==0.26.0 triton==3.6.0 \ - --index-url https://download.pytorch.org/whl/cu130 - -# FlashInfer publishes the CUDA-specific JIT cache on a separate index. -pip install "flashinfer-jit-cache>=0.6.15" \ - --index-url https://flashinfer.ai/whl/cu130 - -pip install -e . +PYTHONNOUSERSITE=1 python -s -m pip install \ + -r requirements/locks/canonical-cu129-py310.txt +PYTHONNOUSERSITE=1 python -s -m pip install --no-deps -e . # Optional MAX_JOBS=8 pip install flash-attn --no-build-isolation @@ -126,44 +121,23 @@ PyTorch wheels include their CUDA runtime, while compiled extensions such as uv venv --python 3.10 source .venv/bin/activate -uv pip install torch==2.11.0 torchvision==0.26.0 triton==3.6.0 \ - --index-url https://download.pytorch.org/whl/cu130 -uv pip install "flashinfer-jit-cache>=0.6.15" \ - --index-url https://flashinfer.ai/whl/cu130 - -uv pip install -e . +uv pip install -r requirements/locks/canonical-cu129-py310.txt +uv pip install --no-deps -e . # Optional MAX_JOBS=8 uv pip install flash-attn --no-build-isolation uv pip install flashinfer-cubin --index-url https://flashinfer.ai/whl ``` -The explicit indexes select the CUDA 13.0 builds of PyTorch and the FlashInfer -JIT cache. - - -Qwen3.5/Qwen3.6 mixed-attention inference additionally requires the -optional Python dependencies: - -```bash -# uv -uv pip install -e ".[qwen35]" - -# Conda/pip -pip install -e ".[qwen35]" -``` - -Vanilla, OmniKV, or QuEST prefix-cache offload without Qwen3.5/Qwen3.6 uses -the smaller CUDA-specific extra: - -```bash -pip install -e ".[prefix-offload]" -``` +The lock selects the CUDA 12.9 builds of PyTorch and the FlashInfer JIT cache. +`einops`, `sgl-kernel`, and the training, benchmark, and test packages are all +part of the main installation; no workflow-specific extras are required. +Sparse-vLLM supports Qwen3.5/Qwen3.6 checkpoints in unquantized BF16 and +block-scaled FP8 formats. The Qwen3.5/Qwen3.6 prefill causal Conv1D and decode Conv1D/GDN packing paths -use repository-local Triton kernels; `sglang-kernel` and a local CUDA extension -build are not required. +use repository-local Triton kernels; they do not call `sgl-kernel` themselves. For the full dependency list and a minimal `LLM(...)` example, see [Getting Started](docs/en/getting_started/README.md). diff --git a/docs/en/getting_started/README.md b/docs/en/getting_started/README.md index d279147b..966db352 100644 --- a/docs/en/getting_started/README.md +++ b/docs/en/getting_started/README.md @@ -10,14 +10,9 @@ Sparse-vLLM usage example. conda create -n svllm python=3.10 -y conda activate svllm -pip install torch==2.11.0 torchvision==0.26.0 triton==3.6.0 \ - --index-url https://download.pytorch.org/whl/cu130 - -# FlashInfer publishes the CUDA-specific JIT cache on a separate index. -pip install "flashinfer-jit-cache>=0.6.15" \ - --index-url https://flashinfer.ai/whl/cu130 - -pip install -e . +PYTHONNOUSERSITE=1 python -s -m pip install \ + -r requirements/locks/canonical-cu129-py310.txt +PYTHONNOUSERSITE=1 python -s -m pip install --no-deps -e . # Optional MAX_JOBS=8 pip install flash-attn --no-build-isolation @@ -25,47 +20,28 @@ MAX_JOBS=8 pip install flash-attn --no-build-isolation ## Install with uv -The project uses the CUDA 13.0 build: +The canonical environment uses the CUDA 12.9 build: ```bash -uv venv --python 3.12 +uv venv --python 3.10 source .venv/bin/activate -uv pip install torch==2.11.0 torchvision==0.26.0 triton==3.6.0 \ - --index-url https://download.pytorch.org/whl/cu130 -uv pip install "flashinfer-jit-cache>=0.6.15" \ - --index-url https://flashinfer.ai/whl/cu130 - -uv pip install -e . +uv pip install -r requirements/locks/canonical-cu129-py310.txt +uv pip install --no-deps -e . # Optional MAX_JOBS=8 uv pip install flash-attn --no-build-isolation ``` -For Qwen3.5/Qwen3.6 mixed-attention inference, install the optional Python -dependencies as well: - -```bash -# uv -uv pip install -e ".[qwen35]" - -# Conda/pip -pip install -e ".[qwen35]" -``` - -For prefix-cache offload with vanilla, OmniKV, or QuEST only: - -```bash -pip install -e ".[prefix-offload]" -``` +`einops`, `sgl-kernel`, and the training, benchmark, and test packages are all +runtime dependencies, so workflow-specific extras are not required. Sparse-vLLM supports Qwen3.5/Qwen3.6 checkpoints in unquantized BF16 and block-scaled FP8 formats. Its prefill causal Conv1D and decode Conv1D/GDN packing paths use local Triton -kernels. They do not require `sglang-kernel` or a repository CUDA-extension -build. +kernels and do not call `sgl-kernel` themselves. The required `flashinfer-jit-cache` package provides modules built for a specific CUDA toolkit version. Select the index matching the CUDA version used @@ -73,8 +49,8 @@ by PyTorch. `flashinfer-cubin` is an optional acceleration package containing architecture-specific device binaries: ```bash -pip install "flashinfer-jit-cache>=0.6.15" \ - --index-url https://flashinfer.ai/whl/cu130 +pip install "flashinfer-jit-cache==0.6.15.post1" \ + --index-url https://flashinfer.ai/whl/cu129 # Optional pip install flashinfer-cubin --index-url https://flashinfer.ai/whl diff --git a/docs/en/getting_started/reproducibility.md b/docs/en/getting_started/reproducibility.md index cc28a28e..1f2204c5 100644 --- a/docs/en/getting_started/reproducibility.md +++ b/docs/en/getting_started/reproducibility.md @@ -9,19 +9,19 @@ evidence, cite the original run artifact path. The README contains the current install command. The expected baseline is: - Python 3.10. -- PyTorch 2.11.0 with the CUDA 13.0 wheel. -- Triton 3.6.0 and torchvision 0.26.0. -- `flashinfer-jit-cache>=0.6.15` installed from the FlashInfer wheel index that - matches `torch.version.cuda`. +- The complete runtime and test environment from + `requirements/locks/canonical-cu129-py310.txt`. +- PyTorch 2.9.1 with the CUDA 12.9 wheel and Triton 3.5.1. +- `flashinfer-python==0.6.15.post1` and the CUDA 12.9 build of + `flashinfer-jit-cache==0.6.15.post1`. +- `sgl-kernel==0.3.21` and `einops>=0.8.2` as runtime dependencies. - Optional matching `flashinfer-cubin` installed from the generic FlashInfer wheel index when precompiled device binaries are desired. -- `transformers[torch]==5.13.1`. +- `transformers==5.13.1`. - `flash-attn` installed with `MAX_JOBS=8 pip install flash-attn --no-build-isolation`. -- Editable install from the repository root with `pip install -e .`. -- Qwen3.5/Qwen3.6 FP8 runs install the CUDA-specific extra with - `pip install -e ".[qwen35]"`. -- Vanilla, OmniKV, and QuEST prefix-cache offload runs install - `pip install -e ".[prefix-offload]"` when the Qwen3.5 extra is not needed. +- Editable install from the repository root with `pip install --no-deps -e .` + after installing the lock. Training, benchmark, and test dependencies are + included in the main installation. - Record the selected operator provider and CUDA compute capability. FP8 providers are selected locally from device capabilities and do not download Hub kernels during warmup. diff --git a/docs/zh/getting_started/README.md b/docs/zh/getting_started/README.md index fe17d7f6..15bbaaa0 100644 --- a/docs/zh/getting_started/README.md +++ b/docs/zh/getting_started/README.md @@ -8,14 +8,9 @@ conda create -n svllm python=3.10 -y conda activate svllm -pip install torch==2.11.0 torchvision==0.26.0 triton==3.6.0 \ - --index-url https://download.pytorch.org/whl/cu130 - -# FlashInfer publishes the CUDA-specific JIT cache on a separate index. -pip install "flashinfer-jit-cache>=0.6.15" \ - --index-url https://flashinfer.ai/whl/cu130 - -pip install -e . +PYTHONNOUSERSITE=1 python -s -m pip install \ + -r requirements/locks/canonical-cu129-py310.txt +PYTHONNOUSERSITE=1 python -s -m pip install --no-deps -e . # Optional MAX_JOBS=8 pip install flash-attn --no-build-isolation @@ -23,48 +18,32 @@ MAX_JOBS=8 pip install flash-attn --no-build-isolation ## 使用 uv 安装 -项目使用 CUDA 13.0 build： +Canonical 环境使用 CUDA 12.9 build： ```bash -uv venv --python 3.12 +uv venv --python 3.10 source .venv/bin/activate -uv pip install torch==2.11.0 torchvision==0.26.0 triton==3.6.0 \ - --index-url https://download.pytorch.org/whl/cu130 -uv pip install "flashinfer-jit-cache>=0.6.15" \ - --index-url https://flashinfer.ai/whl/cu130 - -uv pip install -e . +uv pip install -r requirements/locks/canonical-cu129-py310.txt +uv pip install --no-deps -e . # Optional MAX_JOBS=8 uv pip install flash-attn --no-build-isolation ``` -对于 Qwen3.5/Qwen3.6 mixed-attention inference，还需安装可选 Python 依赖： - -```bash -# uv -uv pip install -e ".[qwen35]" - -# Conda/pip -pip install -e ".[qwen35]" -``` - -如果只需为 vanilla、OmniKV 或 QuEST 启用 prefix-cache offload： - -```bash -pip install -e ".[prefix-offload]" -``` +`einops`、`sgl-kernel` 以及训练、benchmark 和测试包均已是主依赖， +不再需要工作流专用 extra。 Sparse-vLLM 当前支持未量化 BF16 和 block-scaled FP8 格式的 Qwen3.5/Qwen3.6 checkpoint。 -其 prefill causal Conv1D 和 decode Conv1D/GDN packing path 使用仓库本地 Triton kernel，不需要 `sglang-kernel` 或编译仓库 CUDA extension。 +其 prefill causal Conv1D 和 decode Conv1D/GDN packing path 使用仓库本地 +Triton kernel，本身不调用 `sgl-kernel`。 必需的 `flashinfer-jit-cache` package 提供针对特定 CUDA toolkit 版本构建的 module。请选择与 PyTorch 所用 CUDA 版本匹配的 index。`flashinfer-cubin` 是可选加速 package，包含特定架构的 device binary： ```bash -pip install "flashinfer-jit-cache>=0.6.15" \ - --index-url https://flashinfer.ai/whl/cu130 +pip install "flashinfer-jit-cache==0.6.15.post1" \ + --index-url https://flashinfer.ai/whl/cu129 # Optional pip install flashinfer-cubin --index-url https://flashinfer.ai/whl diff --git a/docs/zh/getting_started/reproducibility.md b/docs/zh/getting_started/reproducibility.md index 57a1a3f7..685e2891 100644 --- a/docs/zh/getting_started/reproducibility.md +++ b/docs/zh/getting_started/reproducibility.md @@ -7,15 +7,17 @@ README 包含当前安装命令。预期 baseline 为： - Python 3.10。 -- 带 CUDA 13.0 wheel 的 PyTorch 2.11.0。 -- Triton 3.6.0 和 torchvision 0.26.0。 -- 从与 `torch.version.cuda` 匹配的 FlashInfer wheel index 安装 `flashinfer-jit-cache>=0.6.15`。 +- 使用 `requirements/locks/canonical-cu129-py310.txt` 中冻结的完整 runtime + 与 test 环境。 +- 带 CUDA 12.9 wheel 的 PyTorch 2.9.1，以及 Triton 3.5.1。 +- `flashinfer-python==0.6.15.post1`，以及 CUDA 12.9 build 的 + `flashinfer-jit-cache==0.6.15.post1`。 +- `sgl-kernel==0.3.21` 和 `einops>=0.8.2` 是 runtime 依赖。 - 需要预编译 device binary 时，可从通用 FlashInfer wheel index 安装匹配的 `flashinfer-cubin`。 -- `transformers[torch]==5.13.1`。 +- `transformers==5.13.1`。 - 使用 `MAX_JOBS=8 pip install flash-attn --no-build-isolation` 安装 `flash-attn`。 -- 在仓库根目录运行 `pip install -e .` 进行 editable install。 -- Qwen3.5/Qwen3.6 FP8 run 使用 `pip install -e ".[qwen35]"` 安装 CUDA-specific extra。 -- 不需要 Qwen3.5 extra 时，vanilla、OmniKV 和 QuEST prefix-cache offload run 安装 `pip install -e ".[prefix-offload]"`。 +- 安装 lock 后，在仓库根目录运行 `pip install --no-deps -e .`。训练、 + benchmark 和测试依赖均包含在主安装中。 - 记录选择的 operator provider 和 CUDA compute capability。FP8 provider 根据本地 device capability 选择，warmup 期间不会下载 Hub kernel。 - RMSNorm 默认使用 `SPARSEVLLM_RMSNORM_PROVIDER=auto`，在已安装时优先选择 FlashInfer。设为 `triton` 可强制使用本地 Triton kernel；设为 `flashinfer` 可明确要求 FlashInfer。 diff --git a/pyproject.toml b/pyproject.toml index 72b3a663..b185ac09 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,40 +14,35 @@ authors = [ ] dependencies = [ "fire", - "torch==2.11.0", - "transformers==5.13.1", - "triton==3.6.0", - "torchvision==0.26.0", + "torch>=2.9,<3", + "transformers>=5.13,<6", + "socksio>=1,<2", + "triton>=3.5,<4", + "tilelang==0.1.9", + "apache-tvm-ffi==0.1.10", "pillow", - "datasets", - "matplotlib", - "seaborn", - "loguru", - "ansible", - "bitsandbytes", + "einops", + "sgl-kernel>=0.3.21,<0.4", "tqdm", + "loguru", "fastapi>=0.100", "uvicorn[standard]", "pydantic>=2", - "pytest", + "flashinfer-python>=0.6.15,<0.7", + "flashinfer-jit-cache>=0.6.15,<0.7", + "accelerate", + "datasets", + "wandb", + "bitsandbytes", + "datatrove", + "matplotlib", + "seaborn", "math-verify==0.9.0", - "flashinfer-python>=0.6.15", - "flashinfer-jit-cache>=0.6.15", -] - -[project.optional-dependencies] -prefix-offload = [ - "sgl-kernel @ https://github.com/sgl-project/whl/releases/download/v0.3.14.post1/sgl_kernel-0.3.14.post1%2Bcu128-cp310-abi3-manylinux2014_x86_64.whl", -] -qwen35 = [ - "einops", -] - -test = [ "fuzzywuzzy", "jieba", "pytest", "rouge", + "tomli; python_version < '3.11'", ] [project.scripts] diff --git a/requirements/locks/README.md b/requirements/locks/README.md new file mode 100644 index 00000000..0e79c61a --- /dev/null +++ b/requirements/locks/README.md @@ -0,0 +1,29 @@ +# Dependency locks + +`canonical-cu129-py310.txt` is the supported runtime and test baseline for +Python 3.10 and CUDA 12.9. It freezes the complete resolved Python environment; +`pyproject.toml` describes intentionally broader direct-dependency compatibility +ranges. Training, benchmark, and test packages are part of the main install. +The lock remains the only fully validated package combination; versions inside +the wider ranges are candidates until they pass the same validation gates. + +Create the environment on shared machines under `/data2`: + +```bash +CONDA_PKGS_DIRS=/data2/$USER/cache/conda_pkgs \ + conda create -p /data2/$USER/conda_envs/sparse-vllm-cu129-py310 \ + python=3.10 pip setuptools wheel -y + +PYTHONNOUSERSITE=1 PIP_CACHE_DIR=/data2/$USER/cache/pip \ + /data2/$USER/conda_envs/sparse-vllm-cu129-py310/bin/python -s -m pip \ + install -r requirements/locks/canonical-cu129-py310.txt + +PYTHONNOUSERSITE=1 \ + /data2/$USER/conda_envs/sparse-vllm-cu129-py310/bin/python -s -m pip \ + install --no-deps -e . +``` + +The lock is the compatibility contract. An unlocked install is a development +convenience, not evidence that a newly resolved dependency combination is +supported. Replace the lock only after dependency checks, CPU tests, focused +GPU operator tests, and a real model-path smoke all pass. diff --git a/requirements/locks/canonical-cu129-py310.txt b/requirements/locks/canonical-cu129-py310.txt new file mode 100644 index 00000000..4c2a234e --- /dev/null +++ b/requirements/locks/canonical-cu129-py310.txt @@ -0,0 +1,149 @@ +# Canonical main-dependency lock for Python 3.10 and CUDA 12.9. +# Validated on 2026-08-08 with GLM-4.7-Flash TP=2 and CUDA Graph enabled. +# Keep the CUDA-specific local versions; install from the listed indexes. +--index-url https://pypi.mirrors.ustc.edu.cn/simple +--extra-index-url https://download.pytorch.org/whl/cu129 +--extra-index-url https://flashinfer.ai/whl/cu129 + +accelerate==1.14.0 +aiohappyeyeballs==2.7.1 +aiohttp==3.14.3 +aiosignal==1.4.0 +annotated-doc==0.0.5 +annotated-types==0.8.0 +antlr4-python3-runtime==4.13.2 +anyio==4.14.2 +apache-tvm-ffi==0.1.10 +async-timeout==5.0.1 +attrs==26.1.0 +backports.strenum==1.3.1 +bitsandbytes==0.50.0 +certifi==2026.7.22 +charset-normalizer==3.4.9 +click==8.4.2 +cloudpickle==3.1.2 +contourpy==1.3.2 +cuda-bindings==13.3.1 +cuda-core==1.0.1 +cuda-pathfinder==1.6.0 +cuda-python==13.3.1 +cuda-tile==1.5.0 +cycler==0.12.1 +datasets==4.1.0 +datatrove==0.8.0 +dill==0.4.0 +einops==0.8.2 +exceptiongroup==1.3.1 +fastapi==0.138.0 +filelock==3.29.0 +fire==0.7.1 +flashinfer-jit-cache==0.6.15.post1+cu129 +flashinfer-python==0.6.15.post1 +fonttools==4.63.0 +frozenlist==1.8.0 +fsspec==2025.9.0 +fuzzywuzzy==0.18.0 +h11==0.16.0 +hf-xet==1.6.0 +httpcore==1.0.9 +httptools==0.8.0 +httpx==0.28.1 +huggingface_hub==1.26.1 +humanize==4.16.0 +idna==3.18 +iniconfig==2.3.0 +jieba==0.42.1 +Jinja2==3.1.6 +kiwisolver==1.5.0 +latex2sympy2_extended==1.11.0 +loguru==0.7.3 +markdown-it-py==4.2.0 +MarkupSafe==3.0.3 +math-verify==0.9.0 +matplotlib==3.10.9 +mdurl==0.1.2 +ml_dtypes==0.5.4 +mpmath==1.3.0 +multidict==6.7.1 +multiprocess==0.70.16 +nccl4py==0.3.1 +networkx==3.4.2 +ninja==1.13.0 +numpy==2.2.6 +nvidia-cublas-cu12==12.9.1.4 +nvidia-cuda-cupti-cu12==12.9.79 +nvidia-cuda-nvdisasm==13.3.73 +nvidia-cuda-nvrtc-cu12==12.9.86 +nvidia-cuda-runtime-cu12==12.9.79 +nvidia-cudnn-cu12==9.10.2.21 +nvidia-cudnn-frontend==1.27.0 +nvidia-cufft-cu12==11.4.1.4 +nvidia-cufile-cu12==1.14.1.1 +nvidia-curand-cu12==10.3.10.19 +nvidia-cusolver-cu12==11.7.5.82 +nvidia-cusparse-cu12==12.5.10.65 +nvidia-cusparselt-cu12==0.7.1 +nvidia-cutlass-dsl==4.7.0 +nvidia-cutlass-dsl-libs-base==4.7.0 +nvidia-cutlass-dsl-libs-core==4.7.0 +nvidia-cutlass-dsl-libs-cu12==4.7.0 +nvidia-ml-py==13.610.43 +nvidia-nccl-cu12==2.27.5 +nvidia-nvjitlink-cu12==12.9.86 +nvidia-nvshmem-cu12==3.3.20 +nvidia-nvtx-cu12==12.9.79 +packaging==26.1 +pandas==2.3.3 +pillow==12.1.1 +platformdirs==4.11.0 +pluggy==1.6.0 +propcache==0.5.2 +protobuf==6.33.6 +psutil==7.2.2 +pyarrow==23.0.1 +pydantic==2.12.5 +pydantic_core==2.41.5 +Pygments==2.20.0 +pyparsing==3.3.2 +pytest==9.1.1 +python-dateutil==2.9.0.post0 +python-dotenv==1.2.2 +pytz==2026.3.post1 +PyYAML==6.0.3 +regex==2026.7.19 +requests==2.34.2 +rich==15.0.0 +rouge==1.0.1 +safetensors==0.8.0 +seaborn==0.13.2 +sentry-sdk==2.66.1 +setuptools==80.10.2 +sgl-kernel==0.3.21 +shellingham==1.5.4 +six==1.17.0 +starlette==1.4.1 +sympy==1.14.0 +tabulate==0.10.0 +termcolor==3.3.0 +tilelang==0.1.9 +tokenizers==0.22.2 +tomli==2.4.1 +torch==2.9.1+cu129 +torch_c_dlpack_ext==0.1.5 +tqdm==4.70.0 +transformers==5.13.1 +triton==3.5.1 +typer==0.27.1 +typing-inspection==0.4.2 +typing_extensions==4.15.0 +tzdata==2026.3 +urllib3==2.7.0 +uvicorn==0.49.0 +uvloop==0.22.1 +wandb==0.28.1 +watchfiles==1.2.0 +websockets==16.1.1 +wheel==0.46.3 +xxhash==3.8.1 +yarl==1.24.5 +z3-solver==4.15.4.0 diff --git a/src/sparsevllm/operators/mla_attention.py b/src/sparsevllm/operators/mla_attention.py index d6f79987..b658fdd5 100644 --- a/src/sparsevllm/operators/mla_attention.py +++ b/src/sparsevllm/operators/mla_attention.py @@ -17,6 +17,10 @@ SupportResult, ) from sparsevllm.operators.sgl_fa3 import SglFa3DecodeKernel, sgl_fa3_support +from sparsevllm.operators.tilelang_mla import ( + TileMlaDecodeKernel, + tilelang_mla_support, +) from sparsevllm.platforms.interface import DeviceCaps, PlatformEnum from sparsevllm.triton_kernel.mla import ( DEFAULT_GLM_MLA_DECODE_CONFIG, @@ -566,6 +570,150 @@ def run_explicit_prefill( validation_scope=validation_scope, ) + +@MLA_ATTENTION_REGISTRY.register +class MlaTileLangScoreProvider(MlaSglFa3Provider): + """FA3 output path plus fused TileLang output-and-score for GLM MLA.""" + + name = "tilelang_score_sgl_fa3_h100" + priority = 300 + + def __init__( + self, + *, + op_spec: MlaAttentionOpSpec, + device: torch.device | str, + max_batch_size: int, + launch_config: MlaDecodeLaunchConfig | None = None, + ) -> None: + super().__init__( + op_spec=op_spec, + device=device, + max_batch_size=max_batch_size, + launch_config=launch_config, + ) + self.tilelang_score = TileMlaDecodeKernel( + device=self.device, + softmax_scale=self.spec.softmax_scale, + valid_heads=self.spec.local_q_heads, + ) + + @classmethod + def supports( + cls, + spec: MlaAttentionOpSpec, + caps: DeviceCaps, + ) -> SupportResult: + base = MlaSglFa3Provider.supports(spec, caps) + if not base.supported: + return base + supported, reason = tilelang_mla_support() + return SupportResult.yes(reason) if supported else SupportResult.no(reason) + + @staticmethod + def _tilelang_score_shape_supported( + attn_score: torch.Tensor, + *, + max_context_len: int | None, + ) -> bool: + score_capacity = int(attn_score.shape[1]) if attn_score.ndim >= 2 else 0 + return ( + attn_score.ndim == 2 + and attn_score.dtype == torch.float32 + and score_capacity > 0 + and score_capacity % 64 == 0 + and max_context_len is not None + and int(max_context_len) <= score_capacity + ) + + @staticmethod + def _tilelang_layout_supported( + q_nope_absorbed: torch.Tensor, + q_rope: torch.Tensor, + view: DecodeComputeView, + output: torch.Tensor, + ) -> bool: + if not isinstance(view.payload, MlaLatentPayload): + return False + tensors = ( + q_nope_absorbed, + q_rope, + view.payload.latent_cache, + view.payload.rope_cache, + view.meta.active_slots, + view.meta.req_indices, + view.meta.context_lens, + view.meta.attn_score, + output, + ) + return all( + isinstance(tensor, torch.Tensor) and tensor.is_contiguous() + for tensor in tensors + ) + + @torch.no_grad() + def run( + self, + q_nope_absorbed: torch.Tensor, + q_rope: torch.Tensor, + view: DecodeComputeView, + output: torch.Tensor, + *, + validation_scope: object | None = None, + ) -> torch.Tensor: + attn_score = view.meta.attn_score + if attn_score is None: + return super().run( + q_nope_absorbed, + q_rope, + view, + output, + validation_scope=validation_scope, + ) + # Per-head or non-tile-aligned score buffers remain on the existing + # Triton implementation. This is a static shape dispatch before any + # TileLang kernel launch, not an exception-driven runtime fallback. + if not self._tilelang_score_shape_supported( + attn_score, + max_context_len=view.meta.max_context_len, + ) or not self._tilelang_layout_supported( + q_nope_absorbed, + q_rope, + view, + output, + ): + return MlaTritonProvider.run( + self, + q_nope_absorbed, + q_rope, + view, + output, + validation_scope=validation_scope, + ) + payload = self._validate_run_inputs( + q_nope_absorbed, + q_rope, + view, + output, + ) + self._validate_metadata( + view, + payload, + validation_scope=validation_scope, + ) + return self.tilelang_score( + q_nope_absorbed, + q_rope, + payload.latent_cache, + payload.rope_cache, + view.meta.active_slots, + view.meta.req_indices, + view.meta.context_lens, + output, + attn_score=attn_score, + max_context_len=int(view.meta.max_context_len), + ) + def resolve_mla_attention_provider( spec: MlaAttentionOpSpec, *, @@ -593,6 +741,7 @@ def resolve_mla_attention_provider( "MlaAttentionOpSpec", "MlaAttentionProvider", "MlaSglFa3Provider", + "MlaTileLangScoreProvider", "MlaTritonProvider", "resolve_mla_attention_provider", ] diff --git a/src/sparsevllm/operators/tilelang_mla.py b/src/sparsevllm/operators/tilelang_mla.py new file mode 100644 index 00000000..c37a3da9 --- /dev/null +++ b/src/sparsevllm/operators/tilelang_mla.py @@ -0,0 +1,516 @@ +"""Lazy TileLang adapter for GLM TP1/TP2/TP4 MLA decode. + +The repository-owned TileLang kernel is shape-specialized. This adapter keeps +compilation, padded-query storage, and split-KV workspaces outside the kernel +call and caches them by the CUDA Graph's static batch/context shape. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from importlib import metadata +from typing import Callable + +import torch + + +_VALIDATED_TILELANG_VERSION = "0.1.9" +_VALIDATED_TVM_FFI_VERSION = "0.1.10" +_VALID_SPLITS = (1, 2, 4, 8, 16, 32) +_SUPPORTED_VALID_HEADS = (5, 10, 20) +_SCORE_MODES = ("direct", "atomic", "partial") +_HEAD_TILE_SIZE = 16 +_LATENT_DIM = 512 +_ROPE_DIM = 64 +_BLOCK_N = 64 +_CALIBRATED_CONTEXT_BUCKETS = (1024, 4096, 8192, 16384, 32768, 65536) +_CALIBRATED_BATCH_BUCKETS = (1, 8, 32) +_FUSED_SCORE_SPLITS = { + 5: { + 1: (16, 32, 32, 32, 32, 32), + 8: (16, 16, 16, 16, 16, 32), + 32: (4, 4, 4, 8, 8, 16), + }, + 10: { + 1: (16, 32, 32, 32, 32, 32), + 8: (16, 16, 16, 16, 16, 32), + 32: (4, 4, 4, 8, 4, 4), + }, + 20: { + 1: (16, 32, 32, 32, 32, 32), + 8: (16, 16, 16, 16, 16, 16), + 32: (4, 4, 4, 4, 8, 8), + }, +} +_OUTPUT_ONLY_SPLITS = { + 5: { + 1: (16, 32, 32, 32, 32, 32), + 8: (16, 16, 16, 16, 16, 32), + 32: (4, 4, 4, 8, 4, 4), + }, + 10: { + 1: (16, 32, 32, 32, 32, 32), + 8: (16, 16, 16, 16, 16, 32), + 32: (4, 4, 4, 8, 8, 4), + }, + 20: { + 1: (16, 32, 32, 32, 32, 32), + 8: (16, 16, 16, 16, 16, 16), + 32: (4, 4, 4, 4, 4, 8), + }, +} + + +def _padded_head_count(valid_heads: int) -> int: + if valid_heads not in _SUPPORTED_VALID_HEADS: + raise ValueError( + "TileLang MLA valid_heads must be one of " + f"{_SUPPORTED_VALID_HEADS}, got {valid_heads}." + ) + return 32 if valid_heads > _HEAD_TILE_SIZE else _HEAD_TILE_SIZE + + +def tilelang_mla_support() -> tuple[bool, str]: + """Check the optional package without importing or initializing TileLang.""" + + try: + version = metadata.version("tilelang") + except metadata.PackageNotFoundError: + return False, "tilelang is not installed" + if version != _VALIDATED_TILELANG_VERSION: + return False, ( + f"requires validated tilelang=={_VALIDATED_TILELANG_VERSION}, got " + f"{version!r}" + ) + try: + tvm_ffi_version = metadata.version("apache-tvm-ffi") + except metadata.PackageNotFoundError: + return False, "apache-tvm-ffi is not installed" + if tvm_ffi_version != _VALIDATED_TVM_FFI_VERSION: + return False, ( + "requires validated apache-tvm-ffi==" + f"{_VALIDATED_TVM_FFI_VERSION}, got {tvm_ffi_version!r}" + ) + return True, f"tilelang {version}, apache-tvm-ffi {tvm_ffi_version}" + + +@dataclass(frozen=True, slots=True) +class TileMlaLaunchConfig: + num_split: int + block_n: int = _BLOCK_N + block_h: int = _HEAD_TILE_SIZE + score_mode: str = "direct" + + def __post_init__(self) -> None: + if self.num_split not in _VALID_SPLITS: + raise ValueError( + f"TileLang MLA num_split must be one of {_VALID_SPLITS}, " + f"got {self.num_split}." + ) + if self.block_h not in (16, 32): + raise ValueError( + f"TileLang MLA block_h must be 16 or 32, got {self.block_h}." + ) + if self.score_mode not in _SCORE_MODES: + raise ValueError( + f"TileLang MLA score_mode must be one of {_SCORE_MODES}, " + f"got {self.score_mode!r}." + ) + + +def select_tile_mla_config( + *, + batch_size: int, + context_capacity: int, + need_score: bool, + local_q_heads: int = 10, +) -> TileMlaLaunchConfig: + """Select an offline-calibrated split; never benchmark in the hot path. + + The table is the GPU4 H100 sweep over BS 1/8/32 and context + 1K/4K/8K/16K/32K/64K. Shapes between calibration points use the next + larger bucket; larger shapes reuse the largest calibrated bucket. + """ + + if batch_size <= 0 or context_capacity <= 0: + raise ValueError( + "TileLang MLA batch/context must be positive, got " + f"batch={batch_size} context={context_capacity}." + ) + batch_bucket = next( + (bucket for bucket in _CALIBRATED_BATCH_BUCKETS if batch_size <= bucket), + _CALIBRATED_BATCH_BUCKETS[-1], + ) + context_index = next( + ( + index + for index, bucket in enumerate(_CALIBRATED_CONTEXT_BUCKETS) + if context_capacity <= bucket + ), + len(_CALIBRATED_CONTEXT_BUCKETS) - 1, + ) + _padded_head_count(local_q_heads) + table = _FUSED_SCORE_SPLITS if need_score else _OUTPUT_ONLY_SPLITS + split = table[local_q_heads][batch_bucket][context_index] + block_h = 32 if local_q_heads == 20 and batch_bucket > 1 else 16 + score_mode = "direct" + if need_score and local_q_heads == 20 and block_h == 16: + score_mode = "atomic" if context_capacity <= 4096 else "partial" + return TileMlaLaunchConfig( + num_split=split, + block_h=block_h, + score_mode=score_mode, + ) + + +@dataclass(slots=True) +class TileMlaWorkspace: + padded_latent: torch.Tensor + padded_rope: torch.Tensor + glse: torch.Tensor + partial_output: torch.Tensor + score: torch.Tensor + + +@dataclass(frozen=True, slots=True) +class _KernelKey: + batch_size: int + cache_slot_count: int + active_slot_rows: int + active_slot_width: int + score_capacity: int + num_split: int + block_h: int + score_mode: str + need_score: bool + + +@dataclass(slots=True) +class _BoundKernel: + call: Callable[..., object] + workspace: TileMlaWorkspace + + +class TileMlaDecodeKernel: + """Shape-cached GLM TP1/TP2/TP4 TileLang MLA runner.""" + + def __init__( + self, + *, + device: torch.device | str, + softmax_scale: float, + valid_heads: int = 10, + fixed_config: TileMlaLaunchConfig | None = None, + ) -> None: + self.device = torch.device(device) + self.softmax_scale = float(softmax_scale) + self.valid_heads = int(valid_heads) + self.padded_heads = _padded_head_count(self.valid_heads) + if fixed_config is not None: + if self.padded_heads % fixed_config.block_h: + raise ValueError( + "TileLang MLA padded heads must be divisible by block_h: " + f"padded_heads={self.padded_heads} " + f"block_h={fixed_config.block_h}." + ) + self.fixed_config = fixed_config + self._kernels: dict[_KernelKey, _BoundKernel] = {} + + def _config_for( + self, + *, + batch_size: int, + context_capacity: int, + need_score: bool, + ) -> TileMlaLaunchConfig: + if self.fixed_config is not None: + return self.fixed_config + return select_tile_mla_config( + batch_size=batch_size, + context_capacity=context_capacity, + need_score=need_score, + local_q_heads=self.valid_heads, + ) + + def _bind(self, key: _KernelKey) -> _BoundKernel: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "TileLang MLA shape was not warmed before CUDA Graph capture: " + f"{key}." + ) + supported, reason = tilelang_mla_support() + if not supported: + raise RuntimeError(reason) + # Importing TileLang can initialize its compiler, so keep it behind the + # selected provider and outside module import/resolver paths. + from sparsevllm.tilelang_kernel.mla.decode import ( + build_glm_mla_decode_kernel, + ) + + config = TileMlaLaunchConfig( + key.num_split, + block_h=key.block_h, + score_mode=key.score_mode, + ) + kernel = build_glm_mla_decode_kernel( + batch=key.batch_size, + h_q=self.padded_heads, + h_kv=1, + valid_output_heads=self.valid_heads, + cache_slots=key.cache_slot_count, + slot_rows=key.active_slot_rows, + active_slot_width=key.active_slot_width, + max_seqlen_pad=key.score_capacity, + dv=_LATENT_DIM, + dpe=_ROPE_DIM, + block_N=config.block_n, + block_H=config.block_h, + num_split=config.num_split, + block_size=config.block_n, + softmax_scale=self.softmax_scale, + need_score=key.need_score, + score_mode=config.score_mode, + ) + dtype = torch.bfloat16 + workspace = TileMlaWorkspace( + padded_latent=torch.empty( + key.batch_size, + self.padded_heads, + _LATENT_DIM, + dtype=dtype, + device=self.device, + ), + padded_rope=torch.empty( + key.batch_size, + self.padded_heads, + _ROPE_DIM, + dtype=dtype, + device=self.device, + ), + glse=torch.empty( + key.batch_size, + self.padded_heads, + key.num_split, + dtype=dtype, + device=self.device, + ), + partial_output=torch.empty( + key.batch_size, + self.padded_heads, + key.num_split, + _LATENT_DIM, + dtype=dtype, + device=self.device, + ), + score=torch.empty( + key.batch_size, + self.padded_heads // config.block_h + if config.score_mode == "partial" + else 1, + key.score_capacity, + dtype=torch.float32, + device=self.device, + ), + ) + return _BoundKernel(call=kernel, workspace=workspace) + + def _validate( + self, + q_latent: torch.Tensor, + q_rope: torch.Tensor, + latent_cache: torch.Tensor, + rope_cache: torch.Tensor, + active_slots: torch.Tensor, + request_indices: torch.Tensor, + context_lens: torch.Tensor, + output: torch.Tensor, + attn_score: torch.Tensor | None, + max_context_len: int, + ) -> tuple[int, int]: + batch_size = int(q_latent.shape[0]) + expected = { + "q_latent": (batch_size, self.valid_heads, _LATENT_DIM), + "q_rope": (batch_size, self.valid_heads, _ROPE_DIM), + "output": (batch_size, self.valid_heads, _LATENT_DIM), + "latent_cache": (int(latent_cache.shape[0]), 1, _LATENT_DIM), + "rope_cache": (int(latent_cache.shape[0]), 1, _ROPE_DIM), + "request_indices": (batch_size,), + "context_lens": (batch_size,), + } + actual = { + "q_latent": tuple(q_latent.shape), + "q_rope": tuple(q_rope.shape), + "output": tuple(output.shape), + "latent_cache": tuple(latent_cache.shape), + "rope_cache": tuple(rope_cache.shape), + "request_indices": tuple(request_indices.shape), + "context_lens": tuple(context_lens.shape), + } + for name, shape in expected.items(): + if actual[name] != shape: + raise ValueError( + f"TileLang MLA {name} must have shape {shape}, " + f"got {actual[name]}." + ) + if active_slots.ndim != 2: + raise ValueError( + "TileLang MLA active_slots must be 2D, got " + f"{tuple(active_slots.shape)}." + ) + tensors = { + "q_latent": (q_latent, torch.bfloat16), + "q_rope": (q_rope, torch.bfloat16), + "latent_cache": (latent_cache, torch.bfloat16), + "rope_cache": (rope_cache, torch.bfloat16), + "active_slots": (active_slots, torch.int32), + "request_indices": (request_indices, torch.int32), + "context_lens": (context_lens, torch.int32), + "output": (output, torch.bfloat16), + } + for name, (tensor, dtype) in tensors.items(): + if tensor.device != q_latent.device or tensor.dtype != dtype: + raise TypeError( + f"TileLang MLA {name} must be {dtype} on " + f"{q_latent.device}, got {tensor.dtype} on {tensor.device}." + ) + if not tensor.is_contiguous(): + raise ValueError( + f"TileLang MLA {name} must be contiguous, got stride " + f"{tuple(tensor.stride())}." + ) + score_capacity = int(active_slots.shape[1]) + if attn_score is not None: + if attn_score.ndim != 2 or int(attn_score.shape[0]) != batch_size: + raise ValueError( + "TileLang MLA reduced attn_score must have shape " + f"[batch, capacity], got {tuple(attn_score.shape)}." + ) + if ( + attn_score.dtype != torch.float32 + or attn_score.device != q_latent.device + ): + raise TypeError( + "TileLang MLA attn_score must be FP32 on the query device, " + f"got {attn_score.dtype} on {attn_score.device}." + ) + if not attn_score.is_contiguous(): + raise ValueError( + "TileLang MLA attn_score must be contiguous, got stride " + f"{tuple(attn_score.stride())}." + ) + score_capacity = int(attn_score.shape[1]) + if score_capacity <= 0 or score_capacity % _BLOCK_N: + raise ValueError( + "TileLang MLA context/score capacity must be a positive " + f"multiple of {_BLOCK_N}, got {score_capacity}." + ) + if score_capacity > int(active_slots.shape[1]): + raise ValueError( + "TileLang MLA score capacity exceeds active slot width: " + f"score={score_capacity} slots={active_slots.shape[1]}." + ) + if not 0 < int(max_context_len) <= score_capacity: + raise ValueError( + "TileLang MLA max_context_len must fit the context/score " + f"capacity, got max={max_context_len} capacity={score_capacity}." + ) + return batch_size, score_capacity + + @torch.no_grad() + def __call__( + self, + q_latent: torch.Tensor, + q_rope: torch.Tensor, + latent_cache: torch.Tensor, + rope_cache: torch.Tensor, + active_slots: torch.Tensor, + request_indices: torch.Tensor, + context_lens: torch.Tensor, + output: torch.Tensor, + *, + attn_score: torch.Tensor | None, + max_context_len: int, + ) -> torch.Tensor: + batch_size, score_capacity = self._validate( + q_latent, + q_rope, + latent_cache, + rope_cache, + active_slots, + request_indices, + context_lens, + output, + attn_score, + max_context_len, + ) + need_score = attn_score is not None + config = self._config_for( + batch_size=batch_size, + context_capacity=int(max_context_len), + need_score=need_score, + ) + key = _KernelKey( + batch_size=batch_size, + cache_slot_count=int(latent_cache.shape[0]), + active_slot_rows=int(active_slots.shape[0]), + active_slot_width=int(active_slots.shape[1]), + score_capacity=score_capacity, + num_split=config.num_split, + block_h=config.block_h, + score_mode=config.score_mode, + need_score=need_score, + ) + bound = self._kernels.get(key) + if bound is None: + bound = self._bind(key) + self._kernels[key] = bound + + from sparsevllm.tilelang_kernel.mla.decode import pad_glm_q_kernel + import triton + + workspace = bound.workspace + pad_glm_q_kernel[ + (triton.cdiv(batch_size * self.padded_heads * _LATENT_DIM, 256),) + ]( + q_latent, + q_rope, + workspace.padded_latent, + workspace.padded_rope, + batch_size=batch_size, + valid_heads=self.valid_heads, + padded_heads=self.padded_heads, + latent_dim=_LATENT_DIM, + rope_dim=_ROPE_DIM, + BLOCK=256, + ) + score_output = workspace.score + if attn_score is not None: + if config.score_mode == "partial": + score_output.fill_(-1e20) + else: + score_output = attn_score.unsqueeze(1) + bound.call( + workspace.padded_latent, + workspace.padded_rope, + latent_cache, + rope_cache, + active_slots, + request_indices, + context_lens, + workspace.glse, + workspace.partial_output, + output, + score_output, + ) + if attn_score is not None and config.score_mode == "partial": + torch.amax(score_output, dim=1, out=attn_score) + return output + + +__all__ = [ + "TileMlaDecodeKernel", + "TileMlaLaunchConfig", + "TileMlaWorkspace", + "select_tile_mla_config", + "tilelang_mla_support", +] diff --git a/src/sparsevllm/tilelang_kernel/__init__.py b/src/sparsevllm/tilelang_kernel/__init__.py new file mode 100644 index 00000000..4efc19f9 --- /dev/null +++ b/src/sparsevllm/tilelang_kernel/__init__.py @@ -0,0 +1 @@ +"""Repository-owned adapters for optional TileLang kernels.""" diff --git a/src/sparsevllm/tilelang_kernel/mla/LICENSE.tilelang b/src/sparsevllm/tilelang_kernel/mla/LICENSE.tilelang new file mode 100644 index 00000000..eef5bb23 --- /dev/null +++ b/src/sparsevllm/tilelang_kernel/mla/LICENSE.tilelang @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2024 Tile-AI + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/src/sparsevllm/tilelang_kernel/mla/README.md b/src/sparsevllm/tilelang_kernel/mla/README.md new file mode 100644 index 00000000..7a067cd6 --- /dev/null +++ b/src/sparsevllm/tilelang_kernel/mla/README.md @@ -0,0 +1,28 @@ +# GLM MLA TileLang kernel + +The decode kernel is adapted from +`examples/deepseek_mla/example_mla_decode_paged.py` in Tile-AI/TileLang commit +`c7fabc4cc65e480b88b7606eb1bc9c340dbd8c8c` under the MIT license. + +Local changes implement the GLM-4.7-Flash TP1/TP2/TP4 decode contracts used by +Sparse-vLLM: + +- BF16 query and cache tensors; +- 20/10/5 TP-local query heads padded to complete MMA tiles; +- page-size-one indirect cache slots; +- explicit caller-owned output and split-KV workspaces; +- optional fused FP32 raw-QK score reduced by max over the real local heads + before applying the attention softmax scale; +- indirect request rows and safe `-1` padding outside each context; +- CUDA Graph-compatible execution. + +The module only defines kernels. Provider selection, workspace ownership, +dependency checks, launch-config selection, and fallback policy belong under +`sparsevllm.operators`. + +The production adapter binds the validated GLM TP1/TP2/TP4 H100 BF16 contract. +It chooses an offline-calibrated split, head-tile size, and score reduction +mode from the static batch/context table. TileLang tensors must be contiguous, +and the reduced score/context capacity must be a multiple of the kernel's +64-token tile. Unsupported score dtype, layout, or capacity stays on the +existing Triton provider through an explicit pre-launch shape dispatch. diff --git a/src/sparsevllm/tilelang_kernel/mla/__init__.py b/src/sparsevllm/tilelang_kernel/mla/__init__.py new file mode 100644 index 00000000..b6120d30 --- /dev/null +++ b/src/sparsevllm/tilelang_kernel/mla/__init__.py @@ -0,0 +1,5 @@ +"""GLM MLA TileLang kernels.""" + +from .decode import build_glm_mla_decode_kernel, pad_glm_q_kernel + +__all__ = ["build_glm_mla_decode_kernel", "pad_glm_q_kernel"] diff --git a/src/sparsevllm/tilelang_kernel/mla/decode.py b/src/sparsevllm/tilelang_kernel/mla/decode.py new file mode 100644 index 00000000..566c808b --- /dev/null +++ b/src/sparsevllm/tilelang_kernel/mla/decode.py @@ -0,0 +1,406 @@ +"""TileLang GLM MLA decode kernel. + +Adapted from Tile-AI/TileLang at commit +``c7fabc4cc65e480b88b7606eb1bc9c340dbd8c8c``. The local adaptation uses +BF16, page-size-one indirect cache slots, GLM TP-local query heads padded to +complete MMA tiles, caller-owned outputs/workspaces, and CUDA Graph capture. +See ``LICENSE.tilelang`` and ``README.md`` in this package. +""" + +import tilelang +import triton +import triton.language as tl +import tilelang.language as T + + +@triton.jit +def pad_glm_q_kernel( + q_latent, + q_rope, + padded_latent, + padded_rope, + batch_size: tl.constexpr, + valid_heads: tl.constexpr, + padded_heads: tl.constexpr, + latent_dim: tl.constexpr, + rope_dim: tl.constexpr, + BLOCK: tl.constexpr, +): + offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + latent_total = batch_size * padded_heads * latent_dim + latent_mask = offsets < latent_total + latent_head = (offsets // latent_dim) % padded_heads + latent_batch = offsets // (padded_heads * latent_dim) + latent_col = offsets % latent_dim + latent_src = ( + latent_batch * valid_heads * latent_dim + + latent_head * latent_dim + + latent_col + ) + latent_value = tl.load( + q_latent + latent_src, + mask=latent_mask & (latent_head < valid_heads), + other=0.0, + ) + tl.store(padded_latent + offsets, latent_value, mask=latent_mask) + + rope_total = batch_size * padded_heads * rope_dim + rope_mask = offsets < rope_total + rope_head = (offsets // rope_dim) % padded_heads + rope_batch = offsets // (padded_heads * rope_dim) + rope_col = offsets % rope_dim + rope_src = ( + rope_batch * valid_heads * rope_dim + + rope_head * rope_dim + + rope_col + ) + rope_value = tl.load( + q_rope + rope_src, + mask=rope_mask & (rope_head < valid_heads), + other=0.0, + ) + tl.store(padded_rope + offsets, rope_value, mask=rope_mask) + + +@tilelang.jit( + out_idx=[], + pass_configs={ + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + }, +) +def build_glm_mla_decode_kernel( + batch, + h_q, + h_kv, + valid_output_heads, + cache_slots, + slot_rows, + active_slot_width, + max_seqlen_pad, + dv, + dpe, + block_N, + block_H, + num_split, + block_size, + softmax_scale=None, + need_score=False, + score_mode="direct", +): + if softmax_scale is None: + softmax_scale = (dv + dpe) ** -0.5 + scale = float(softmax_scale * 1.44269504) # log2(e) + dtype = T.bfloat16 + accum_dtype = T.float32 + kv_group_num = h_q // h_kv + VALID_BLOCK_H = min(block_H, kv_group_num) + VALID_OUTPUT_HEADS = valid_output_heads + HEAD_TILE_COUNT = h_q // VALID_BLOCK_H + SCORE_TILE_COUNT = HEAD_TILE_COUNT if score_mode == "partial" else 1 + assert h_kv == 1, "h_kv must be 1" + assert h_q % VALID_BLOCK_H == 0, "h_q must use complete head tiles" + assert 0 < VALID_OUTPUT_HEADS <= h_q, "valid output heads must fit h_q" + assert score_mode in ("direct", "atomic", "partial") + assert not need_score or score_mode != "direct" or HEAD_TILE_COUNT == 1 + assert block_size >= block_N and block_size % block_N == 0, ( + "block_size must be at least block_N and a multiple of block_N" + ) + + @T.prim_func + def main_split( + Q: T.Tensor([batch, h_q, dv], dtype), + Q_pe: T.Tensor([batch, h_q, dpe], dtype), + KV: T.Tensor([cache_slots, h_kv, dv], dtype), + K_pe: T.Tensor([cache_slots, h_kv, dpe], dtype), + active_slots: T.Tensor([slot_rows, active_slot_width], T.int32), + request_indices: T.Tensor([batch], T.int32), + cache_seqlens: T.Tensor([batch], T.int32), + glse: T.Tensor([batch, h_q, num_split], dtype), + Output_partial: T.Tensor([batch, h_q, num_split, dv], dtype), + Output: T.Tensor([batch, VALID_OUTPUT_HEADS, dv], dtype), + AttnScore: T.Tensor( + [batch, SCORE_TILE_COUNT, max_seqlen_pad], accum_dtype + ), + ): + # split kv + with T.Kernel(batch, h_q // min(block_H, kv_group_num), num_split, threads=256) as (bx, by, bz): + Q_shared = T.alloc_shared([block_H, dv], dtype) + S_shared = T.alloc_shared([block_H, block_N], dtype) + Q_pe_shared = T.alloc_shared([block_H, dpe], dtype) + KV_shared = T.alloc_shared([block_N, dv], dtype) + K_pe_shared = T.alloc_shared([block_N, dpe], dtype) + O_shared = T.alloc_shared([block_H, dv], dtype) + acc_s = T.alloc_fragment([block_H, block_N], accum_dtype) + acc_s_cast = T.alloc_fragment([block_H, block_N], dtype) + acc_o = T.alloc_fragment([block_H, dv], accum_dtype) + scores_max = T.alloc_fragment([block_H], accum_dtype) + scores_max_prev = T.alloc_fragment([block_H], accum_dtype) + scores_scale = T.alloc_fragment([block_H], accum_dtype) + scores_sum = T.alloc_fragment([block_H], accum_dtype) + token_scores = T.alloc_fragment([block_N], accum_dtype) + logsum = T.alloc_fragment([block_H], accum_dtype) + + cur_kv_head = 0 + request_row = T.max(request_indices[bx], 0) + if HEAD_TILE_COUNT == 1: + T.use_swizzle(10) + + T.copy(Q[bx, by * VALID_BLOCK_H : (by + 1) * VALID_BLOCK_H, :], Q_shared) + T.copy(Q_pe[bx, by * VALID_BLOCK_H : (by + 1) * VALID_BLOCK_H, :], Q_pe_shared) + T.fill(acc_o, 0) + T.fill(logsum, 0) + T.fill(scores_max, -T.infinity(accum_dtype)) + + total_blocks = T.ceildiv(cache_seqlens[bx], block_N) + blocks_per_split = T.floordiv(total_blocks, num_split) + remaining_blocks = T.floormod(total_blocks, num_split) + loop_range = blocks_per_split + T.if_then_else(bz < remaining_blocks, 1, 0) + start = (blocks_per_split * bz + T.min(bz, remaining_blocks)) * block_N + + for k in T.Pipelined(loop_range, num_stages=2): + for i, j in T.Parallel(block_N, dv): + token_index = start + k * block_N + i + slot = T.if_then_else( + token_index < cache_seqlens[bx], + active_slots[request_row, token_index], + 0, + ) + KV_shared[i, j] = KV[slot, cur_kv_head, j] + for i, j in T.Parallel(block_N, dpe): + token_index = start + k * block_N + i + slot = T.if_then_else( + token_index < cache_seqlens[bx], + active_slots[request_row, token_index], + 0, + ) + K_pe_shared[i, j] = K_pe[slot, cur_kv_head, j] + T.clear(acc_s) + T.gemm(Q_shared, KV_shared, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullCol) + T.gemm(Q_pe_shared, K_pe_shared, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullCol) + T.copy(scores_max, scores_max_prev) + T.fill(scores_max, -T.infinity(accum_dtype)) + for i, j in T.Parallel(block_H, block_N): + acc_s[i, j] = T.if_then_else(start + k * block_N + j >= cache_seqlens[bx], -T.infinity(accum_dtype), acc_s[i, j]) + for i, j in T.Parallel(block_H, block_N): + acc_s[i, j] = T.if_then_else(by * VALID_BLOCK_H + i >= VALID_OUTPUT_HEADS, -T.infinity(accum_dtype), acc_s[i, j]) + if need_score: + T.reduce_max(acc_s, token_scores, dim=0) + if score_mode == "direct": + for j in T.Parallel(block_N): + score_index = start + k * block_N + j + AttnScore[bx, 0, score_index] = T.if_then_else( + score_index < cache_seqlens[bx], + token_scores[j], + AttnScore[bx, 0, score_index], + ) + elif score_mode == "atomic": + for j in T.Parallel(block_N): + score_index = start + k * block_N + j + if score_index < cache_seqlens[bx]: + T.atomic_max( + AttnScore[bx, 0, score_index], + token_scores[j], + ) + else: + for j in T.Parallel(block_N): + score_index = start + k * block_N + j + AttnScore[bx, by, score_index] = T.if_then_else( + score_index < cache_seqlens[bx], + token_scores[j], + AttnScore[bx, by, score_index], + ) + T.reduce_max(acc_s, scores_max, dim=1, clear=False) + for i in T.Parallel(block_H): + scores_max[i] = T.max(scores_max[i], scores_max_prev[i]) + for i in T.Parallel(block_H): + scores_scale[i] = T.exp2(scores_max_prev[i] * scale - scores_max[i] * scale) + for i, j in T.Parallel(block_H, block_N): + acc_s[i, j] = T.exp2(acc_s[i, j] * scale - scores_max[i] * scale) + T.reduce_sum(acc_s, scores_sum, dim=1) + T.copy(acc_s, S_shared) + T.copy(S_shared, acc_s_cast) + for i in T.Parallel(block_H): + logsum[i] = logsum[i] * scores_scale[i] + scores_sum[i] + for i, j in T.Parallel(block_H, dv): + acc_o[i, j] *= scores_scale[i] + T.gemm(acc_s_cast, KV_shared, acc_o, policy=T.GemmWarpPolicy.FullCol) + for i, j in T.Parallel(block_H, dv): + acc_o[i, j] = T.if_then_else( + loop_range > 0, + acc_o[i, j] / logsum[i], + 0.0, + ) + for i in T.Parallel(block_H): + logsum[i] = T.if_then_else( + loop_range > 0, + T.log2(logsum[i]) + scores_max[i] * scale, + -T.infinity(accum_dtype), + ) + T.copy(logsum, glse[bx, by * VALID_BLOCK_H : (by + 1) * VALID_BLOCK_H, bz]) + T.copy(acc_o, O_shared) + T.copy(O_shared, Output_partial[bx, by * VALID_BLOCK_H : (by + 1) * VALID_BLOCK_H, bz, :]) + + # combine + with T.Kernel(VALID_OUTPUT_HEADS, batch, threads=128) as (by, bz): + po_local = T.alloc_fragment([dv], dtype) + o_accum_local = T.alloc_fragment([dv], accum_dtype) + lse_local_split = T.alloc_var(accum_dtype) + lse_logsum_local = T.alloc_var(accum_dtype) + lse_max_local = T.alloc_var(accum_dtype) + scale_local = T.alloc_var(accum_dtype) + + T.clear(lse_logsum_local) + T.clear(o_accum_local) + lse_max_local = -T.infinity(accum_dtype) + for k in T.serial(num_split): + lse_max_local = T.max(lse_max_local, glse[bz, by, k]) + for k in T.Pipelined(num_split, num_stages=1): + lse_local_split = glse[bz, by, k] + lse_logsum_local += T.exp2(lse_local_split - lse_max_local) + lse_logsum_local = T.log2(lse_logsum_local) + lse_max_local + for k in T.serial(num_split): + for i in T.Parallel(dv): + po_local[i] = Output_partial[bz, by, k, i] + lse_local_split = glse[bz, by, k] + scale_local = T.exp2(lse_local_split - lse_logsum_local) + for i in T.Parallel(dv): + o_accum_local[i] += po_local[i] * scale_local + for i in T.Parallel(dv): + Output[bz, by, i] = T.if_then_else( + cache_seqlens[bz] > 0, + o_accum_local[i], + 0.0, + ) + + @T.prim_func + def main_no_split( + Q: T.Tensor([batch, h_q, dv], dtype), + Q_pe: T.Tensor([batch, h_q, dpe], dtype), + KV: T.Tensor([cache_slots, h_kv, dv], dtype), + K_pe: T.Tensor([cache_slots, h_kv, dpe], dtype), + active_slots: T.Tensor([slot_rows, active_slot_width], T.int32), + request_indices: T.Tensor([batch], T.int32), + cache_seqlens: T.Tensor([batch], T.int32), + glse: T.Tensor([batch, h_q, num_split], dtype), + Output_partial: T.Tensor([batch, h_q, num_split, dv], dtype), + Output: T.Tensor([batch, VALID_OUTPUT_HEADS, dv], dtype), + AttnScore: T.Tensor( + [batch, SCORE_TILE_COUNT, max_seqlen_pad], accum_dtype + ), + ): + with T.Kernel(batch, h_q // min(block_H, kv_group_num), threads=256) as (bx, by): + Q_shared = T.alloc_shared([block_H, dv], dtype) + S_shared = T.alloc_shared([block_H, block_N], dtype) + Q_pe_shared = T.alloc_shared([block_H, dpe], dtype) + KV_shared = T.alloc_shared([block_N, dv], dtype) + K_pe_shared = T.alloc_shared([block_N, dpe], dtype) + O_shared = T.alloc_shared([block_H, dv], dtype) + acc_s = T.alloc_fragment([block_H, block_N], accum_dtype) + acc_o = T.alloc_fragment([block_H, dv], accum_dtype) + scores_max = T.alloc_fragment([block_H], accum_dtype) + scores_max_prev = T.alloc_fragment([block_H], accum_dtype) + scores_scale = T.alloc_fragment([block_H], accum_dtype) + scores_sum = T.alloc_fragment([block_H], accum_dtype) + token_scores = T.alloc_fragment([block_N], accum_dtype) + logsum = T.alloc_fragment([block_H], accum_dtype) + + cur_kv_head = 0 + request_row = T.max(request_indices[bx], 0) + if HEAD_TILE_COUNT == 1: + T.use_swizzle(10) + + T.copy(Q[bx, by * VALID_BLOCK_H : (by + 1) * VALID_BLOCK_H, :], Q_shared) + T.copy(Q_pe[bx, by * VALID_BLOCK_H : (by + 1) * VALID_BLOCK_H, :], Q_pe_shared) + T.fill(acc_o, 0) + T.fill(logsum, 0) + T.fill(scores_max, -T.infinity(accum_dtype)) + + loop_range = T.ceildiv(cache_seqlens[bx], block_N) + for kr in T.Pipelined(loop_range, num_stages=2): + k = loop_range - 1 - kr + for i, j in T.Parallel(block_N, dv): + token_index = k * block_N + i + slot = T.if_then_else( + token_index < cache_seqlens[bx], + active_slots[request_row, token_index], + 0, + ) + KV_shared[i, j] = KV[slot, cur_kv_head, j] + for i, j in T.Parallel(block_N, dpe): + token_index = k * block_N + i + slot = T.if_then_else( + token_index < cache_seqlens[bx], + active_slots[request_row, token_index], + 0, + ) + K_pe_shared[i, j] = K_pe[slot, cur_kv_head, j] + T.clear(acc_s) + T.gemm(Q_shared, KV_shared, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullCol) + T.gemm(Q_pe_shared, K_pe_shared, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullCol) + T.copy(scores_max, scores_max_prev) + T.fill(scores_max, -T.infinity(accum_dtype)) + if kr == 0: + for i, j in T.Parallel(block_H, block_N): + acc_s[i, j] = T.if_then_else(k * block_N + j >= cache_seqlens[bx], -T.infinity(accum_dtype), acc_s[i, j]) + for i, j in T.Parallel(block_H, block_N): + acc_s[i, j] = T.if_then_else(by * VALID_BLOCK_H + i >= VALID_OUTPUT_HEADS, -T.infinity(accum_dtype), acc_s[i, j]) + if need_score: + T.reduce_max(acc_s, token_scores, dim=0) + if score_mode == "direct": + for j in T.Parallel(block_N): + score_index = k * block_N + j + AttnScore[bx, 0, score_index] = T.if_then_else( + score_index < cache_seqlens[bx], + token_scores[j], + AttnScore[bx, 0, score_index], + ) + elif score_mode == "atomic": + for j in T.Parallel(block_N): + score_index = k * block_N + j + if score_index < cache_seqlens[bx]: + T.atomic_max( + AttnScore[bx, 0, score_index], + token_scores[j], + ) + else: + for j in T.Parallel(block_N): + score_index = k * block_N + j + AttnScore[bx, by, score_index] = T.if_then_else( + score_index < cache_seqlens[bx], + token_scores[j], + AttnScore[bx, by, score_index], + ) + T.reduce_max(acc_s, scores_max, dim=1, clear=False) + for i in T.Parallel(block_H): + scores_max[i] = T.max(scores_max[i], scores_max_prev[i]) + for i in T.Parallel(block_H): + scores_scale[i] = T.exp2(scores_max_prev[i] * scale - scores_max[i] * scale) + for i, j in T.Parallel(block_H, block_N): + acc_s[i, j] = T.exp2(acc_s[i, j] * scale - scores_max[i] * scale) + T.reduce_sum(acc_s, scores_sum, dim=1) + T.copy(acc_s, S_shared) + for i in T.Parallel(block_H): + logsum[i] = logsum[i] * scores_scale[i] + scores_sum[i] + for i, j in T.Parallel(block_H, dv): + acc_o[i, j] *= scores_scale[i] + T.gemm(S_shared, KV_shared, acc_o, policy=T.GemmWarpPolicy.FullCol) + for i, j in T.Parallel(block_H, dv): + acc_o[i, j] = T.if_then_else( + cache_seqlens[bx] > 0, + acc_o[i, j] / logsum[i], + 0.0, + ) + T.copy(acc_o, O_shared) + for i, j in T.Parallel(block_H, dv): + global_head = by * VALID_BLOCK_H + i + if global_head < VALID_OUTPUT_HEADS: + Output[bx, global_head, j] = T.if_then_else( + cache_seqlens[bx] > 0, + O_shared[i, j], + 0.0, + ) + + if num_split > 1: + return main_split + else: + return main_no_split diff --git a/src/sparsevllm/utils/select_omnikv_full_layers.py b/src/sparsevllm/utils/select_omnikv_full_layers.py index 435b151d..f7005a72 100644 --- a/src/sparsevllm/utils/select_omnikv_full_layers.py +++ b/src/sparsevllm/utils/select_omnikv_full_layers.py @@ -6,6 +6,7 @@ import random import subprocess import sys +import typing from dataclasses import asdict, dataclass from datetime import datetime from pathlib import Path @@ -66,6 +67,21 @@ def require_path(path: str | Path, kind: str) -> Path: return resolved +def install_typing_compatibility() -> list[str]: + installed: list[str] = [] + if not hasattr(typing, "Unpack"): + try: + from typing_extensions import Unpack + except ImportError as exc: + raise RuntimeError( + "Python < 3.11 requires typing_extensions.Unpack to load " + "the MiniMax Transformers remote model code." + ) from exc + typing.Unpack = Unpack # type: ignore[attr-defined] + installed.append("typing.Unpack") + return installed + + def text_model_config(config): return getattr(config, "text_config", config) @@ -382,6 +398,14 @@ def move_inputs(token_ids: list[int], device: torch.device) -> torch.Tensor: return torch.tensor([token_ids], dtype=torch.long, device=device) +def model_input_device(model, fallback: torch.device) -> torch.device: + embeddings = model.get_input_embeddings() + weight = getattr(embeddings, "weight", None) + if weight is None or weight.device.type == "meta": + return fallback + return weight.device + + @torch.no_grad() def advance_cache(model, past_key_values, token_ids: list[int], *, device: torch.device, chunk_size: int): if chunk_size <= 0: @@ -475,13 +499,20 @@ def run_calibration(args: argparse.Namespace) -> dict[str, Any]: if args.dataset not in dataset2prompt: raise ValueError(f"Dataset {args.dataset!r} is missing from {prompt_path}.") prompt_format = dataset2prompt[args.dataset] + typing_compatibility = ( + install_typing_compatibility() if args.trust_remote_code else [] + ) timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") output_dir = Path(args.output_dir) if args.output_dir else Path(args.output_root) / f"{args.dataset}_{timestamp}" output_dir.mkdir(parents=True, exist_ok=False) - tokenizer = AutoTokenizer.from_pretrained(str(model_path), trust_remote_code=True) - base_config = AutoConfig.from_pretrained(str(model_path), trust_remote_code=True) + tokenizer = AutoTokenizer.from_pretrained( + str(model_path), trust_remote_code=args.trust_remote_code + ) + base_config = AutoConfig.from_pretrained( + str(model_path), trust_remote_code=args.trust_remote_code + ) text_config = text_model_config(base_config) attention_layer_indices = attention_layer_indices_from_config(base_config) max_length = int(args.max_length or getattr(text_config, "max_position_embeddings", 32000)) @@ -489,25 +520,52 @@ def run_calibration(args: argparse.Namespace) -> dict[str, Any]: raise ValueError(f"Resolved max_length must be > 0, got {max_length}.") dtype = torch_dtype_from_name(args.torch_dtype) - device = torch.device(args.device) - if device.type == "cuda" and not torch.cuda.is_available(): + requested_device = torch.device(args.device) + if requested_device.type == "cuda" and not torch.cuda.is_available(): raise RuntimeError("CUDA device requested, but torch.cuda.is_available() is False.") + if args.max_memory_per_device_gib <= 0: + raise ValueError( + "max_memory_per_device_gib must be > 0, " + f"got {args.max_memory_per_device_gib}." + ) removed_fp8_exclusions = prepare_fp8_transformers_config(base_config) model_kwargs = { "config": base_config, "torch_dtype": dtype, - "trust_remote_code": True, + "trust_remote_code": args.trust_remote_code, "attn_implementation": "eager", } - if device.type == "cuda": - model_kwargs["device_map"] = {"": str(device)} + if args.device_map == "auto": + if requested_device.type != "cuda": + raise ValueError("--device-map auto requires a CUDA --device.") + visible_devices = torch.cuda.device_count() + if visible_devices <= 0: + raise RuntimeError("--device-map auto requires at least one visible CUDA device.") + model_kwargs["device_map"] = "auto" + model_kwargs["max_memory"] = { + index: f"{args.max_memory_per_device_gib}GiB" + for index in range(visible_devices) + } + elif requested_device.type == "cuda": + model_kwargs["device_map"] = {"": str(requested_device)} model = AutoModelForCausalLM.from_pretrained( str(model_path), **model_kwargs, ) - if device.type != "cuda": - model.to(device) + if args.device_map == "auto": + hf_device_map = getattr(model, "hf_device_map", {}) + offloaded = sorted( + {str(value) for value in hf_device_map.values() if str(value) in {"cpu", "disk"}} + ) + if offloaded: + raise RuntimeError( + "Automatic device placement offloaded model modules to " + f"{offloaded}; increase visible GPU memory instead of running a mixed CPU/disk calibration." + ) + elif requested_device.type != "cuda": + model.to(requested_device) + device = model_input_device(model, requested_device) model.eval() num_hidden_layers = int(text_config.num_hidden_layers) @@ -635,9 +693,18 @@ def run_calibration(args: argparse.Namespace) -> dict[str, Any]: "model_config_num_hidden_layers": num_hidden_layers, "attention_layer_indices": attention_layer_indices, "removed_fp8_modules_to_not_convert": removed_fp8_exclusions, + "typing_compatibility": typing_compatibility, + "trust_remote_code": bool(args.trust_remote_code), "attention_implementation": "eager", "torch_dtype": args.torch_dtype, "device": args.device, + "device_map": args.device_map, + "max_memory_per_device_gib": int(args.max_memory_per_device_gib), + "resolved_input_device": str(device), + "hf_device_map": { + str(key): str(value) + for key, value in getattr(model, "hf_device_map", {}).items() + }, "prefill_chunk_size": int(args.prefill_chunk_size), "max_length": int(max_length), "no_chat_template": bool(args.no_chat_template), @@ -669,6 +736,13 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: parser.add_argument("--prefill-chunk-size", type=int, default=512) parser.add_argument("--max-length", type=int, default=None) parser.add_argument("--device", default="cuda") + parser.add_argument( + "--trust-remote-code", + action=argparse.BooleanOptionalAction, + default=True, + ) + parser.add_argument("--device-map", default="single", choices=("single", "auto")) + parser.add_argument("--max-memory-per-device-gib", type=int, default=76) parser.add_argument("--torch-dtype", default="bfloat16") parser.add_argument("--thinking-mode", default="off", choices=("off", "on")) parser.add_argument("--no-chat-template", action="store_true") diff --git a/tests/test_dependency_constraints.py b/tests/test_dependency_constraints.py index d3df2451..c7837fe4 100644 --- a/tests/test_dependency_constraints.py +++ b/tests/test_dependency_constraints.py @@ -1,18 +1,70 @@ +from pathlib import Path + try: import tomllib -except ModuleNotFoundError: # Python 3.10, which this project supports. +except ModuleNotFoundError: # Python 3.10 import tomli as tomllib -from pathlib import Path -def test_flashinfer_minimum_version_matches_moe_api(): +def test_runtime_compatibility_bounds_cover_canonical_lock(): pyproject_path = Path(__file__).parents[1] / "pyproject.toml" project = tomllib.loads(pyproject_path.read_text())["project"] dependencies = set(project["dependencies"]) - assert "flashinfer-python>=0.6.15" in dependencies - assert "flashinfer-jit-cache>=0.6.15" in dependencies + assert "torch>=2.9,<3" in dependencies + assert "triton>=3.5,<4" in dependencies + assert "tilelang==0.1.9" in dependencies + assert "apache-tvm-ffi==0.1.10" in dependencies + assert "transformers>=5.13,<6" in dependencies + assert "flashinfer-python>=0.6.15,<0.7" in dependencies + assert "flashinfer-jit-cache>=0.6.15,<0.7" in dependencies + assert "sgl-kernel>=0.3.21,<0.4" in dependencies + assert {"fire", "pillow", "einops", "tqdm", "loguru"} <= dependencies + assert not any( + dependency.startswith("torchvision") for dependency in dependencies + ) assert not any( dependency.startswith("flashinfer-cubin") for dependency in dependencies ) + + +def test_canonical_lock_pins_validated_tilelang_runtime(): + lock_path = ( + Path(__file__).parents[1] + / "requirements" + / "locks" + / "canonical-cu129-py310.txt" + ) + locked_requirements = { + line.strip() + for line in lock_path.read_text().splitlines() + if line and not line.startswith(("#", "--")) + } + + assert "tilelang==0.1.9" in locked_requirements + assert "apache-tvm-ffi==0.1.10" in locked_requirements + + +def test_workflow_dependencies_are_part_of_main_install(): + pyproject_path = Path(__file__).parents[1] / "pyproject.toml" + project = tomllib.loads(pyproject_path.read_text())["project"] + dependencies = set(project["dependencies"]) + + assert "optional-dependencies" not in project + assert { + "accelerate", + "datasets", + "socksio>=1,<2", + "wandb", + "bitsandbytes", + "datatrove", + "matplotlib", + "seaborn", + "math-verify==0.9.0", + "fuzzywuzzy", + "jieba", + "pytest", + "rouge", + "tomli; python_version < '3.11'", + } <= dependencies diff --git a/tests/test_mla_attention_operator.py b/tests/test_mla_attention_operator.py index 87040b2d..b7c2ea51 100644 --- a/tests/test_mla_attention_operator.py +++ b/tests/test_mla_attention_operator.py @@ -15,6 +15,7 @@ MLA_ATTENTION_REGISTRY, MlaAttentionOpSpec, MlaSglFa3Provider, + MlaTileLangScoreProvider, MlaTritonProvider, ) from sparsevllm.operators.registry import OpResolver @@ -118,6 +119,7 @@ def test_mla_resolver_selects_h100_triton_provider() -> None: assert isinstance(resolved.provider, MlaTritonProvider) assert resolved.rejected == ( ("sgl_fa3_h100", "sgl-kernel is not installed"), + ("tilelang_score_sgl_fa3_h100", "sgl-kernel is not installed"), ) allocate.assert_called_once_with( batch_size=8, @@ -136,6 +138,10 @@ def test_mla_resolver_prefers_sgl_fa3_when_available() -> None: "sparsevllm.operators.mla_attention.sgl_fa3_support", return_value=(True, "validated test API"), ), + patch( + "sparsevllm.operators.mla_attention.tilelang_mla_support", + return_value=(False, "tilelang unavailable"), + ), patch( "sparsevllm.operators.mla_attention.allocate_mla_decode_workspace", return_value=workspace, @@ -151,16 +157,33 @@ def test_mla_resolver_prefers_sgl_fa3_when_available() -> None: ) assert type(resolved.provider) is MlaSglFa3Provider - assert resolved.rejected == () + assert resolved.rejected == ( + ( + "tilelang_score_sgl_fa3_h100", + "tilelang unavailable", + ), + ) def test_mla_resolver_accepts_cuda_graph_after_capture_gate() -> None: spec = _spec(cuda_graph=True) workspace = _cpu_workspace(batch_size=8, head_count=spec.local_q_heads) - with patch( - "sparsevllm.operators.mla_attention.allocate_mla_decode_workspace", - return_value=workspace, + with ( + patch( + "sparsevllm.operators.mla_attention.sgl_fa3_support", + return_value=(True, "validated test API"), + ), + patch( + "sparsevllm.operators.mla_attention.tilelang_mla_support", + return_value=(True, "validated test API"), + ), + patch( + "sparsevllm.operators.mla_attention.allocate_mla_decode_workspace", + return_value=workspace, + ), + patch("sparsevllm.operators.mla_attention.SglFa3DecodeKernel"), + patch("sparsevllm.operators.mla_attention.TileMlaDecodeKernel"), ): resolved = OpResolver(MLA_ATTENTION_REGISTRY).resolve( spec, @@ -170,7 +193,7 @@ def test_mla_resolver_accepts_cuda_graph_after_capture_gate() -> None: max_batch_size=8, ) - assert isinstance(resolved.provider, MlaTritonProvider) + assert type(resolved.provider) is MlaTileLangScoreProvider assert resolved.rejected == () diff --git a/tests/test_omnikv_full_layer_selector.py b/tests/test_omnikv_full_layer_selector.py index 6ca3fee7..9ca270f5 100644 --- a/tests/test_omnikv_full_layer_selector.py +++ b/tests/test_omnikv_full_layer_selector.py @@ -1,14 +1,19 @@ import random +import typing import unittest from types import SimpleNamespace import numpy as np +import torch from sparsevllm.utils.select_omnikv_full_layers import ( CalibrationPoint, add_topk_to_pair_scores, attention_layer_indices_from_config, compute_segment_scores, + install_typing_compatibility, + model_input_device, + parse_args, prepare_fp8_transformers_config, sample_decode_points, select_full_layers_dp, @@ -17,6 +22,39 @@ class OmniKVFullLayerSelectorTest(unittest.TestCase): + def test_remote_model_code_can_be_disabled(self): + args = parse_args( + [ + "--model-path", + "/model", + "--longbench-root", + "/longbench", + "--no-trust-remote-code", + ] + ) + + self.assertFalse(args.trust_remote_code) + + def test_python310_typing_compatibility_is_idempotent(self): + had_unpack = hasattr(typing, "Unpack") + + first = install_typing_compatibility() + second = install_typing_compatibility() + + self.assertTrue(hasattr(typing, "Unpack")) + self.assertEqual(second, []) + self.assertEqual(first, [] if had_unpack else ["typing.Unpack"]) + + def test_model_input_device_uses_embedding_weight_device(self): + class FakeModel: + def get_input_embeddings(self): + return SimpleNamespace(weight=torch.empty(1, device="cpu")) + + self.assertEqual( + model_input_device(FakeModel(), torch.device("cuda:0")), + torch.device("cpu"), + ) + def test_add_topk_to_pair_scores_counts_forward_layer_intersections(self): pair_scores = np.zeros((4, 4), dtype=np.int64) add_topk_to_pair_scores( diff --git a/tests/test_tilelang_mla_kernel.py b/tests/test_tilelang_mla_kernel.py new file mode 100644 index 00000000..5055be57 --- /dev/null +++ b/tests/test_tilelang_mla_kernel.py @@ -0,0 +1,205 @@ +from __future__ import annotations + +import pytest +import torch + +from sparsevllm.operators.tilelang_mla import ( + TileMlaDecodeKernel, + TileMlaLaunchConfig, +) + + +CUDA_REQUIRED = pytest.mark.skipif( + not torch.cuda.is_available(), + reason="CUDA is required for the TileLang MLA kernel test", +) + + +def _torch_oracle( + q_latent: torch.Tensor, + q_rope: torch.Tensor, + latent_cache: torch.Tensor, + rope_cache: torch.Tensor, + slots: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + latent_keys = latent_cache[slots.long(), 0].float() + rope_keys = rope_cache[slots.long(), 0].float() + raw = torch.matmul(q_latent.float(), latent_keys.T) + torch.matmul( + q_rope.float(), rope_keys.T + ) + score = raw.max(dim=0).values + probability = torch.softmax(raw * (256**-0.5), dim=-1) + output = torch.matmul(probability, latent_keys) + return output.to(torch.bfloat16), score + + +@CUDA_REQUIRED +@pytest.mark.parametrize( + ("valid_heads", "num_split", "block_h", "score_mode"), + [ + (5, 1, 16, "direct"), + (5, 4, 16, "direct"), + (10, 1, 16, "direct"), + (10, 32, 16, "direct"), + (20, 1, 16, "atomic"), + (20, 4, 16, "partial"), + (20, 4, 32, "direct"), + ], +) +def test_tilelang_mla_score_matches_torch_with_indirect_slots_and_graph( + valid_heads: int, + num_split: int, + block_h: int, + score_mode: str, +) -> None: + torch.manual_seed(20260808 + valid_heads + num_split) + device = torch.device("cuda") + batch_size = 2 + capacity = 64 + cache_slots = 96 + q_latent = torch.randn( + batch_size, valid_heads, 512, dtype=torch.bfloat16, device=device + ) + q_rope = torch.randn( + batch_size, valid_heads, 64, dtype=torch.bfloat16, device=device + ) + latent_cache = torch.randn( + cache_slots, 1, 512, dtype=torch.bfloat16, device=device + ) + rope_cache = torch.randn( + cache_slots, 1, 64, dtype=torch.bfloat16, device=device + ) + active_slots = torch.full( + (3, capacity), -1, dtype=torch.int32, device=device + ) + active_slots[0, :17] = torch.arange( + 50, 67, dtype=torch.int32, device=device + ) + active_slots[2, :33] = torch.randperm( + cache_slots, dtype=torch.int64, device=device + )[:33].to(torch.int32) + request_indices = torch.tensor([2, -1], dtype=torch.int32, device=device) + context_lens = torch.tensor([33, 0], dtype=torch.int32, device=device) + output = torch.empty_like(q_latent) + score = torch.full( + (batch_size, capacity), + -1e20, + dtype=torch.float32, + device=device, + ) + runner = TileMlaDecodeKernel( + device=device, + softmax_scale=256**-0.5, + valid_heads=valid_heads, + fixed_config=TileMlaLaunchConfig( + num_split, + block_h=block_h, + score_mode=score_mode, + ), + ) + + def run() -> None: + score.fill_(-1e20) + runner( + q_latent, + q_rope, + latent_cache, + rope_cache, + active_slots, + request_indices, + context_lens, + output, + attn_score=score, + max_context_len=capacity, + ) + + run() + torch.cuda.synchronize() + expected_output, expected_score = _torch_oracle( + q_latent[0], + q_rope[0], + latent_cache, + rope_cache, + active_slots[2, :33], + ) + torch.testing.assert_close( + output[0], expected_output, rtol=3e-2, atol=3e-2 + ) + torch.testing.assert_close( + score[0, :33], expected_score, rtol=3e-2, atol=3e-2 + ) + torch.testing.assert_close(output[1], torch.zeros_like(output[1])) + assert torch.all(score[0, 33:] == -1e20) + assert torch.all(score[1] == -1e20) + + run() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + graph.replay() + torch.cuda.synchronize() + graph_output = output.clone() + graph_score = score.clone() + run() + torch.cuda.synchronize() + torch.testing.assert_close(graph_output, output, rtol=0, atol=0) + torch.testing.assert_close(graph_score, score, rtol=0, atol=0) + + +@CUDA_REQUIRED +@pytest.mark.parametrize("valid_heads", [5, 20]) +def test_tilelang_score_ignores_zero_padded_heads(valid_heads: int) -> None: + device = torch.device("cuda") + capacity = 64 + q_latent = torch.ones( + 1, valid_heads, 512, dtype=torch.bfloat16, device=device + ) + q_rope = torch.ones( + 1, valid_heads, 64, dtype=torch.bfloat16, device=device + ) + latent_cache = -torch.ones( + capacity, 1, 512, dtype=torch.bfloat16, device=device + ) + rope_cache = -torch.ones( + capacity, 1, 64, dtype=torch.bfloat16, device=device + ) + active_slots = torch.arange( + capacity, dtype=torch.int32, device=device + ).unsqueeze(0) + request_indices = torch.zeros(1, dtype=torch.int32, device=device) + context_lens = torch.full((1,), capacity, dtype=torch.int32, device=device) + output = torch.empty_like(q_latent) + score = torch.full( + (1, capacity), -1e20, dtype=torch.float32, device=device + ) + runner = TileMlaDecodeKernel( + device=device, + softmax_scale=256**-0.5, + valid_heads=valid_heads, + fixed_config=TileMlaLaunchConfig( + 1, + block_h=32 if valid_heads == 20 else 16, + ), + ) + + runner( + q_latent, + q_rope, + latent_cache, + rope_cache, + active_slots, + request_indices, + context_lens, + output, + attn_score=score, + max_context_len=capacity, + ) + torch.cuda.synchronize() + + torch.testing.assert_close( + score, + torch.full_like(score, -576.0), + rtol=0, + atol=0, + ) diff --git a/tests/test_tilelang_mla_operator.py b/tests/test_tilelang_mla_operator.py new file mode 100644 index 00000000..e3865f91 --- /dev/null +++ b/tests/test_tilelang_mla_operator.py @@ -0,0 +1,439 @@ +from __future__ import annotations + +from importlib import metadata +from unittest.mock import Mock, patch + +import pytest +import torch + +from sparsevllm.engine.cache_manager import ( + AttentionViewMeta, + DecodeComputeView, + MlaLatentPayload, +) +from sparsevllm.operators.mla_attention import ( + MLA_ATTENTION_REGISTRY, + MlaAttentionOpSpec, + MlaSglFa3Provider, + MlaTileLangScoreProvider, +) +from sparsevllm.operators.registry import OpResolver +from sparsevllm.operators.tilelang_mla import ( + TileMlaDecodeKernel, + select_tile_mla_config, + tilelang_mla_support, +) +from sparsevllm.platforms import DeviceCaps, PlatformEnum +from sparsevllm.triton_kernel.mla import MlaDecodeWorkspace + + +def _spec(*, tp_size: int = 2) -> MlaAttentionOpSpec: + return MlaAttentionOpSpec( + num_q_heads=20, + kv_lora_rank=512, + rope_dim=64, + qk_head_dim=256, + value_head_dim=256, + activation_dtype=torch.bfloat16, + cache_dtype=torch.bfloat16, + tp_size=tp_size, + cuda_graph=True, + ) + + +def _h100_caps() -> DeviceCaps: + return DeviceCaps( + platform=PlatformEnum.CUDA, + device_type="cuda", + device_index=0, + device_name="NVIDIA H100 80GB HBM3", + compute_capability=(9, 0), + runtime_version="13.0", + supports_graph_capture=True, + supports_torch_compile=True, + supports_triton=True, + supports_pin_memory=True, + supports_bfloat16=True, + supports_native_fp8=True, + ) + + +def _cpu_workspace() -> MlaDecodeWorkspace: + return MlaDecodeWorkspace( + block_size=torch.empty(1, dtype=torch.int32), + batch_start_indices=torch.empty(2, dtype=torch.int32), + mid_output=torch.empty(20, 1, 512, dtype=torch.float32), + mid_logsumexp=torch.empty(20, 1, dtype=torch.float32), + ) + + +def _view(*, score: torch.Tensor | None) -> DecodeComputeView: + active_slots = torch.full((3, 64), -1, dtype=torch.int32) + active_slots[2, :3] = torch.tensor([5, 2, 7], dtype=torch.int32) + return DecodeComputeView( + meta=AttentionViewMeta( + active_slots=active_slots, + req_indices=torch.tensor([2, -1], dtype=torch.int32), + context_lens=torch.tensor([3, 0], dtype=torch.int32), + max_context_len=64, + attn_score=score, + ), + payload=MlaLatentPayload( + latent_cache=torch.empty(8, 1, 512, dtype=torch.bfloat16), + rope_cache=torch.empty(8, 1, 64, dtype=torch.bfloat16), + ), + ) + + +def test_tilelang_support_does_not_import_runtime() -> None: + def version(package: str) -> str: + return {"tilelang": "0.1.9", "apache-tvm-ffi": "0.1.10"}[package] + + with patch.object(metadata, "version", side_effect=version): + assert tilelang_mla_support() == ( + True, + "tilelang 0.1.9, apache-tvm-ffi 0.1.10", + ) + + +@pytest.mark.parametrize( + ("version", "supported"), + [("0.1.8", False), ("0.1.9", True), ("0.2.0+cu130", False)], +) +def test_tilelang_support_version_boundary(version: str, supported: bool) -> None: + def installed(package: str) -> str: + return version if package == "tilelang" else "0.1.10" + + with patch.object(metadata, "version", side_effect=installed): + assert tilelang_mla_support()[0] is supported + + +def test_tilelang_support_reports_missing_package() -> None: + with patch.object( + metadata, + "version", + side_effect=metadata.PackageNotFoundError, + ): + assert tilelang_mla_support() == (False, "tilelang is not installed") + + +def test_tilelang_support_rejects_unvalidated_tvm_ffi() -> None: + def version(package: str) -> str: + return { + "tilelang": "0.1.9", + "apache-tvm-ffi": "0.1.13.post2", + }[package] + + with patch.object(metadata, "version", side_effect=version): + supported, reason = tilelang_mla_support() + assert not supported + assert "apache-tvm-ffi==0.1.10" in reason + + +@pytest.mark.parametrize( + ( + "heads", + "batch", + "context", + "need_score", + "expected", + ), + [ + (10, 1, 1024, True, (16, 16, "direct")), + (10, 8, 65536, True, (32, 16, "direct")), + (5, 32, 32768, True, (8, 16, "direct")), + (20, 1, 1024, True, (16, 16, "atomic")), + (20, 1, 4096, True, (32, 16, "atomic")), + (20, 1, 8192, True, (32, 16, "partial")), + (20, 8, 16384, True, (16, 32, "direct")), + (20, 32, 32768, True, (8, 32, "direct")), + (20, 64, 131072, False, (8, 32, "direct")), + ], +) +def test_tilelang_config_selection_is_static( + heads: int, + batch: int, + context: int, + need_score: bool, + expected: tuple[int, int, str], +) -> None: + config = select_tile_mla_config( + batch_size=batch, + context_capacity=context, + need_score=need_score, + local_q_heads=heads, + ) + assert (config.num_split, config.block_h, config.score_mode) == expected + + +@pytest.mark.parametrize("tp_size", [1, 2, 4]) +def test_tilelang_provider_supports_all_glm_tp_sizes(tp_size: int) -> None: + with ( + patch( + "sparsevllm.operators.mla_attention.sgl_fa3_support", + return_value=(True, "sgl test"), + ), + patch( + "sparsevllm.operators.mla_attention.tilelang_mla_support", + return_value=(True, "tilelang test"), + ), + ): + assert MlaTileLangScoreProvider.supports( + _spec(tp_size=tp_size), _h100_caps() + ).supported + + +def test_resolver_prefers_tilelang_score_provider_for_tp2() -> None: + workspace = _cpu_workspace() + with ( + patch( + "sparsevllm.operators.mla_attention.sgl_fa3_support", + return_value=(True, "sgl test"), + ), + patch( + "sparsevllm.operators.mla_attention.tilelang_mla_support", + return_value=(True, "tilelang test"), + ), + patch( + "sparsevllm.operators.mla_attention.allocate_mla_decode_workspace", + return_value=workspace, + ), + patch("sparsevllm.operators.mla_attention.SglFa3DecodeKernel"), + patch("sparsevllm.operators.mla_attention.TileMlaDecodeKernel"), + ): + resolved = OpResolver(MLA_ATTENTION_REGISTRY).resolve( + _spec(), + _h100_caps(), + op_spec=_spec(), + device="cpu", + max_batch_size=2, + ) + assert type(resolved.provider) is MlaTileLangScoreProvider + + +@pytest.mark.parametrize(("tp_size", "local_heads"), [(1, 20), (2, 10), (4, 5)]) +def test_tilelang_provider_binds_rank_local_head_count( + tp_size: int, + local_heads: int, +) -> None: + workspace = _cpu_workspace() + with ( + patch( + "sparsevllm.operators.mla_attention.allocate_mla_decode_workspace", + return_value=workspace, + ), + patch("sparsevllm.operators.mla_attention.SglFa3DecodeKernel"), + patch( + "sparsevllm.operators.mla_attention.TileMlaDecodeKernel" + ) as tilelang_cls, + ): + MlaTileLangScoreProvider( + op_spec=_spec(tp_size=tp_size), + device="cpu", + max_batch_size=2, + ) + + tilelang_cls.assert_called_once_with( + device=torch.device("cpu"), + softmax_scale=256**-0.5, + valid_heads=local_heads, + ) + + +def test_missing_tilelang_keeps_existing_sgl_provider() -> None: + workspace = _cpu_workspace() + with ( + patch( + "sparsevllm.operators.mla_attention.sgl_fa3_support", + return_value=(True, "sgl test"), + ), + patch( + "sparsevllm.operators.mla_attention.tilelang_mla_support", + return_value=(False, "tilelang missing"), + ), + patch( + "sparsevllm.operators.mla_attention.allocate_mla_decode_workspace", + return_value=workspace, + ), + patch("sparsevllm.operators.mla_attention.SglFa3DecodeKernel"), + ): + resolved = OpResolver(MLA_ATTENTION_REGISTRY).resolve( + _spec(), + _h100_caps(), + op_spec=_spec(), + device="cpu", + max_batch_size=2, + ) + assert type(resolved.provider) is MlaSglFa3Provider + assert ( + "tilelang_score_sgl_fa3_h100", + "tilelang missing", + ) in resolved.rejected + + +def _provider_with_mocks() -> tuple[MlaTileLangScoreProvider, Mock, Mock]: + fa3 = Mock(return_value=torch.empty(2, 10, 512, dtype=torch.bfloat16)) + tilelang = Mock(return_value=torch.empty(2, 10, 512, dtype=torch.bfloat16)) + with ( + patch( + "sparsevllm.operators.mla_attention.allocate_mla_decode_workspace", + return_value=_cpu_workspace(), + ), + patch( + "sparsevllm.operators.mla_attention.SglFa3DecodeKernel", + return_value=fa3, + ), + patch( + "sparsevllm.operators.mla_attention.TileMlaDecodeKernel", + return_value=tilelang, + ), + ): + provider = MlaTileLangScoreProvider( + op_spec=_spec(), + device="cpu", + max_batch_size=2, + ) + return provider, fa3, tilelang + + +def test_score_path_routes_to_tilelang_with_caller_owned_score() -> None: + provider, fa3, tilelang = _provider_with_mocks() + score = torch.full((2, 64), -1e20, dtype=torch.float32) + view = _view(score=score) + q_latent = torch.empty(2, 10, 512, dtype=torch.bfloat16) + q_rope = torch.empty(2, 10, 64, dtype=torch.bfloat16) + output = torch.empty_like(q_latent) + + with patch( + "sparsevllm.operators.mla_attention.validate_mla_decode_metadata" + ): + provider.run(q_latent, q_rope, view, output) + + fa3.assert_not_called() + tilelang.assert_called_once_with( + q_latent, + q_rope, + view.payload.latent_cache, + view.payload.rope_cache, + view.meta.active_slots, + view.meta.req_indices, + view.meta.context_lens, + output, + attn_score=score, + max_context_len=64, + ) + + +def test_no_score_path_remains_fa3() -> None: + provider, fa3, tilelang = _provider_with_mocks() + view = _view(score=None) + q_latent = torch.empty(2, 10, 512, dtype=torch.bfloat16) + q_rope = torch.empty(2, 10, 64, dtype=torch.bfloat16) + output = torch.empty_like(q_latent) + + with patch( + "sparsevllm.operators.mla_attention.validate_mla_decode_metadata" + ): + provider.run(q_latent, q_rope, view, output) + + tilelang.assert_not_called() + fa3.assert_called_once() + + +@pytest.mark.parametrize( + "score", + [ + torch.empty(2, 10, 64, dtype=torch.float32), + torch.empty(2, 64, dtype=torch.bfloat16), + torch.empty(2, 63, dtype=torch.float32), + ], +) +def test_unsupported_score_contract_uses_explicit_triton_path(score) -> None: + provider, fa3, tilelang = _provider_with_mocks() + view = _view(score=score) + q_latent = torch.empty(2, 10, 512, dtype=torch.bfloat16) + q_rope = torch.empty(2, 10, 64, dtype=torch.bfloat16) + output = torch.empty_like(q_latent) + + with patch.object( + MlaSglFa3Provider.__mro__[1], "run", return_value=output + ) as triton: + provider.run(q_latent, q_rope, view, output) + + fa3.assert_not_called() + tilelang.assert_not_called() + triton.assert_called_once() + + +def test_score_capacity_smaller_than_declared_context_uses_triton() -> None: + provider, fa3, tilelang = _provider_with_mocks() + view = _view(score=torch.empty(2, 64, dtype=torch.float32)) + object.__setattr__(view.meta, "max_context_len", 128) + q_latent = torch.empty(2, 10, 512, dtype=torch.bfloat16) + q_rope = torch.empty(2, 10, 64, dtype=torch.bfloat16) + output = torch.empty_like(q_latent) + + with patch.object( + MlaSglFa3Provider.__mro__[1], "run", return_value=output + ) as triton: + provider.run(q_latent, q_rope, view, output) + + fa3.assert_not_called() + tilelang.assert_not_called() + triton.assert_called_once() + + +@pytest.mark.parametrize("noncontiguous", ["active_slots", "attn_score"]) +def test_noncontiguous_tilelang_inputs_use_triton(noncontiguous: str) -> None: + provider, fa3, tilelang = _provider_with_mocks() + score = torch.empty(2, 128, dtype=torch.float32)[:, ::2] + view = _view( + score=( + score + if noncontiguous == "attn_score" + else torch.empty(2, 64, dtype=torch.float32) + ) + ) + if noncontiguous == "active_slots": + backing = torch.full((3, 128), -1, dtype=torch.int32) + backing[2, :6:2] = torch.tensor([5, 2, 7], dtype=torch.int32) + object.__setattr__(view.meta, "active_slots", backing[:, ::2]) + q_latent = torch.empty(2, 10, 512, dtype=torch.bfloat16) + q_rope = torch.empty(2, 10, 64, dtype=torch.bfloat16) + output = torch.empty_like(q_latent) + + with patch.object( + MlaSglFa3Provider.__mro__[1], "run", return_value=output + ) as triton: + provider.run(q_latent, q_rope, view, output) + + fa3.assert_not_called() + tilelang.assert_not_called() + triton.assert_called_once() + + +def test_tilelang_runner_rejects_unaligned_score_capacity_before_import() -> None: + runner = TileMlaDecodeKernel(device="cpu", softmax_scale=0.0625) + view = _view(score=torch.empty(2, 63, dtype=torch.float32)) + with pytest.raises(ValueError, match="multiple of 64"): + runner( + torch.empty(2, 10, 512, dtype=torch.bfloat16), + torch.empty(2, 10, 64, dtype=torch.bfloat16), + view.payload.latent_cache, + view.payload.rope_cache, + view.meta.active_slots, + view.meta.req_indices, + view.meta.context_lens, + torch.empty(2, 10, 512, dtype=torch.bfloat16), + attn_score=view.meta.attn_score, + max_context_len=63, + ) + + +def test_tilelang_runner_rejects_unsupported_local_head_count() -> None: + with pytest.raises(ValueError, match="valid_heads must be one of"): + TileMlaDecodeKernel( + device="cpu", + softmax_scale=0.0625, + valid_heads=7, + ) From 58dad01cc4040e15e7056630e743df7067814167 Mon Sep 17 00:00:00 2001 From: CURRENTF <1206013760@qq.com> Date: Sun, 9 Aug 2026 05:42:50 +0800 Subject: [PATCH 19/48] perf: optimize glm tp2 ep2 runtime --- src/sparsevllm/layers/mla_attention.py | 13 +++++ src/sparsevllm/models/glm4_moe_lite.py | 21 +++++-- src/sparsevllm/operators/mla_attention.py | 37 +++++++++++- src/sparsevllm/operators/sgl_fa3.py | 53 +++++++++++++++++ src/sparsevllm/triton_kernel/moe_config.py | 34 ++++++++++- tests/test_glm4_moe_lite.py | 21 +++++++ tests/test_mla_attention_layer.py | 9 +++ tests/test_mla_attention_operator.py | 66 +++++++++++++++++++++- tests/test_moe_config.py | 36 ++++++++++++ tests/test_sgl_fa3.py | 26 ++++++++- tests/test_tilelang_mla_operator.py | 1 + 11 files changed, 308 insertions(+), 9 deletions(-) diff --git a/src/sparsevllm/layers/mla_attention.py b/src/sparsevllm/layers/mla_attention.py index 9da76fa5..eda27f45 100644 --- a/src/sparsevllm/layers/mla_attention.py +++ b/src/sparsevllm/layers/mla_attention.py @@ -33,6 +33,7 @@ class MlaPrefillHistory: gathered_latent: torch.Tensor gathered_rope: torch.Tensor packed_offsets: torch.Tensor + packed_cu_seqlens: torch.Tensor packed_slots: torch.Tensor local_req_indices: torch.Tensor context_lens: torch.Tensor @@ -66,6 +67,7 @@ class _MlaPrefillPlan: source_query_tokens: int cache_slot_count: int packed_offsets: torch.Tensor + packed_cu_seqlens: torch.Tensor packed_slots: torch.Tensor local_req_indices: torch.Tensor context_lengths: tuple[int, ...] @@ -469,6 +471,11 @@ def _get_prefill_plan( dtype=torch.int32, device=self.device, ) + packed_cu_seqlens = torch.tensor( + (*packed_starts, total_visible_tokens), + dtype=torch.int32, + device=self.device, + ) local_req_indices = torch.arange( batch_size, dtype=torch.int32, @@ -498,6 +505,7 @@ def _get_prefill_plan( source_query_tokens=int(query_tokens), cache_slot_count=int(cache_slot_count), packed_offsets=packed_offsets, + packed_cu_seqlens=packed_cu_seqlens, packed_slots=packed_slots, local_req_indices=local_req_indices, context_lengths=lengths, @@ -554,6 +562,7 @@ def prepare_prefill_history( gathered_latent=gathered_latent, gathered_rope=gathered_rope, packed_offsets=plan.packed_offsets, + packed_cu_seqlens=plan.packed_cu_seqlens, packed_slots=plan.packed_slots, local_req_indices=plan.local_req_indices, context_lens=meta.context_lens, @@ -832,6 +841,10 @@ def build_prefill_explicit_view( payload=ExplicitKVPayload( k_cache=workset.expanded_k, v_cache=workset.expanded_v, + metadata={ + "layout": "mla_packed_varlen", + "cu_seqlens_k": history.packed_cu_seqlens, + }, ), ) diff --git a/src/sparsevllm/models/glm4_moe_lite.py b/src/sparsevllm/models/glm4_moe_lite.py index 873bf941..0a4fa191 100644 --- a/src/sparsevllm/models/glm4_moe_lite.py +++ b/src/sparsevllm/models/glm4_moe_lite.py @@ -577,7 +577,7 @@ def __init__( hidden_act=str(config.hidden_act), mlp_chunk_size=self.mlp_chunk_size, quantization=None, - reduce_results=self.parallel_context.ep_size > 1, + reduce_results=False, ) ) @@ -686,16 +686,27 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: ) if debug_enabled: # Preserve the general EP composition and routed-only debug - # evidence. The single-reduction specialization below is TP-only. + # evidence while keeping shared-expert reductions explicit. self.parallel_context.world_all_reduce(routed) self.debug_last_routed_output = routed.detach().clone() shared = self._shared_chunk(hidden_states) - if self.parallel_context.ep_size == 1: + if self.parallel_context.tp_size > 1: shared = self.parallel_context.tp_all_reduce(shared) output = routed + shared elif self.parallel_context.ep_size > 1: - self.parallel_context.world_all_reduce(routed) - output = routed + self._shared_chunk(hidden_states) + shared_local = self._shared_chunk(hidden_states) + if self.parallel_context.tp_size > 1: + # Hybrid TP+EP makes both branches partial over the same + # outer world. Sum them locally so one collective completes + # routed experts and the TP-sharded shared expert together. + output = self.parallel_context.world_all_reduce_out_of_place( + routed + shared_local + ) + else: + # Retain the pure-EP semantic path for direct module use: the + # shared expert is replicated rather than TP-sharded. + self.parallel_context.world_all_reduce(routed) + output = routed + shared_local else: shared_local = self._shared_chunk(hidden_states) if context.is_prefill: diff --git a/src/sparsevllm/operators/mla_attention.py b/src/sparsevllm/operators/mla_attention.py index b658fdd5..d1770e32 100644 --- a/src/sparsevllm/operators/mla_attention.py +++ b/src/sparsevllm/operators/mla_attention.py @@ -471,7 +471,8 @@ def run( view.meta.req_indices, view.meta.context_lens, output, - num_splits=5, + # Zero enables FA3's measured context-aware split heuristic. + num_splits=0, validation_scope=validation_scope, ) @@ -557,6 +558,40 @@ def run_explicit_prefill( raise TypeError( f"{name} must use {expected_dtype}, got {tensor.dtype}." ) + metadata = payload.metadata or {} + if metadata.get("layout") == "mla_packed_varlen": + cu_seqlens_k = metadata.get("cu_seqlens_k") + if not isinstance(cu_seqlens_k, torch.Tensor): + raise TypeError( + "MLA packed varlen prefill requires tensor cu_seqlens_k." + ) + if cu_seqlens_k.shape != (batch_size + 1,): + raise ValueError( + f"cu_seqlens_k must have shape ({batch_size + 1},), got " + f"{tuple(cu_seqlens_k.shape)}." + ) + if ( + cu_seqlens_k.device != self.device + or cu_seqlens_k.dtype != torch.int32 + ): + raise TypeError( + "cu_seqlens_k must be int32 on the provider device, got " + f"{cu_seqlens_k.device}/{cu_seqlens_k.dtype}." + ) + if view.meta.max_context_len is None: + raise ValueError( + "MLA packed varlen prefill requires max_context_len." + ) + return self.fa3.run_contiguous_explicit_varlen( + q, + payload.k_cache, + payload.v_cache, + output, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=int(max_seqlen_q), + max_seqlen_k=int(view.meta.max_context_len), + ) return self.fa3.run_explicit_varlen( q, payload.k_cache, diff --git a/src/sparsevllm/operators/sgl_fa3.py b/src/sparsevllm/operators/sgl_fa3.py index f029cdfd..e2f31807 100644 --- a/src/sparsevllm/operators/sgl_fa3.py +++ b/src/sparsevllm/operators/sgl_fa3.py @@ -506,5 +506,58 @@ def run_explicit_varlen( raise RuntimeError("sgl-kernel FA3 did not write to the supplied output") return output + def run_contiguous_explicit_varlen( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + output: torch.Tensor, + *, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + ) -> torch.Tensor: + """Run causal varlen attention over packed contiguous KV.""" + + args: list[object] = [ + q, + k, + v, + None, + None, + None, + output, + cu_seqlens_q, + cu_seqlens_k, + None, + None, + None, + int(max_seqlen_q), + int(max_seqlen_k), + None, + None, + None, + None, + None, + None, + None, + None, + None, + self.softmax_scale, + True, + -1, + -1, + ] + if self._has_attention_chunk: + args.append(0) + args.extend((0.0, True, None, self.num_splits, None, 0, None)) + if self._has_04_arguments: + args.extend((None, False)) + result: Sequence[torch.Tensor] = self._op(*args) + if not result or result[0].data_ptr() != output.data_ptr(): + raise RuntimeError("sgl-kernel FA3 did not write to the supplied output") + return output + __all__ = ["SglFa3DecodeKernel", "sgl_fa3_support"] diff --git a/src/sparsevllm/triton_kernel/moe_config.py b/src/sparsevllm/triton_kernel/moe_config.py index 0fe696d1..e7de53d5 100644 --- a/src/sparsevllm/triton_kernel/moe_config.py +++ b/src/sparsevllm/triton_kernel/moe_config.py @@ -100,6 +100,8 @@ def _heuristic_config( _GLM_DECODE_64 = MoeGemmConfig(64, 128, 64, 8, 8, 3) _GLM_MID_BATCH = MoeGemmConfig(64, 128, 64, 1, 8, 3) _GLM_LARGE_BATCH = MoeGemmConfig(128, 128, 64, 1, 8, 3) +_GLM_EP2_TINY_BATCH = MoeGemmConfig(16, 128, 32, 8, 4, 4) +_GLM_EP2_SMALL_BATCH = MoeGemmConfig(16, 64, 128, 1, 4, 4) def _glm_h100_tp2_config( @@ -141,6 +143,32 @@ def _glm_h100_tp2_config( return _GLM_LARGE_BATCH +def _glm_h100_tp2_ep2_config( + shape: MoeGemmShape, + *, + num_tokens: int, + stage: str, +) -> MoeGemmConfig | None: + """Return measured BF16 configs for the GLM outer-TP2/EP2 shape.""" + + profiled_shape = MoeGemmShape( + "NVIDIA H100 80GB HBM3", + (9, 0), + torch.bfloat16, + 4, + 32, + 2048, + 1536, + ) + if stage not in {"w13", "w2"} or shape != profiled_shape: + return None + if num_tokens <= 4: + return _GLM_EP2_TINY_BATCH + if num_tokens <= 128: + return _GLM_EP2_SMALL_BATCH + return _GLM_LARGE_BATCH + + def _stage_table( w13: tuple[MoeGemmConfig, ...], w2: tuple[MoeGemmConfig, ...], @@ -434,7 +462,11 @@ def _resolve_moe_gemm_config( hidden_size=hidden_size, intermediate_size=intermediate_size, ) - glm_config = _glm_h100_tp2_config( + glm_config = _glm_h100_tp2_ep2_config( + shape, + num_tokens=num_tokens, + stage=stage, + ) or _glm_h100_tp2_config( shape, num_tokens=num_tokens, stage=stage, diff --git a/tests/test_glm4_moe_lite.py b/tests/test_glm4_moe_lite.py index 9019ae1c..bc3e29c7 100644 --- a/tests/test_glm4_moe_lite.py +++ b/tests/test_glm4_moe_lite.py @@ -498,6 +498,27 @@ def test_glm_sparse_moe_reduces_hybrid_tp_ep_shards_over_outer_world() -> None: assert all_reduce.call_args.kwargs["group"] is world_process_group +def test_glm_hybrid_tp_ep_shared_expert_defers_reduction_to_moe_block() -> None: + world_process_group = object() + context = ParallelContext( + world=ParallelGroup(world_process_group, (0, 1), 0, 2), + tensor=ParallelGroup(world_process_group, (0, 1), 0, 2), + expert=ParallelGroup(world_process_group, (0, 1), 0, 2), + data=ParallelGroup(None, (0,), 0, 1), + moe_tensor=ParallelGroup(None, (0,), 0, 1), + ) + config = _config() + with _construction_context(context): + block = Glm4MoeLiteSparseMoeBlock( + config, + mlp_chunk_size=config.mlp_chunk_size, + decode_cuda_graph=False, + ) + + assert block.shared_experts is not None + assert block.shared_experts.down_proj.reduce_results is False + + def test_glm_moe_debug_contract_populates_model_runner_summaries() -> None: context = _ep_context() block = object.__new__(Glm4MoeLiteSparseMoeBlock) diff --git a/tests/test_mla_attention_layer.py b/tests/test_mla_attention_layer.py index 49424c95..873de62f 100644 --- a/tests/test_mla_attention_layer.py +++ b/tests/test_mla_attention_layer.py @@ -313,6 +313,7 @@ def test_mla_prefill_reuses_validated_packing_across_layers() -> None: assert validate.call_count == 2 assert gather.call_count == 3 assert first.packed_offsets is second.packed_offsets + assert first.packed_cu_seqlens is second.packed_cu_seqlens assert first.packed_slots is second.packed_slots @@ -385,6 +386,14 @@ def test_mla_prefill_exposes_expanded_explicit_score_view() -> None: assert isinstance(score_view.payload, ExplicitKVPayload) assert score_view.payload.k_cache is workset.expanded_k assert score_view.payload.v_cache is workset.expanded_v + assert score_view.payload.metadata == { + "layout": "mla_packed_varlen", + "cu_seqlens_k": history.packed_cu_seqlens, + } + torch.testing.assert_close( + history.packed_cu_seqlens, + torch.tensor([0, 2], dtype=torch.int32), + ) assert score_view.meta.active_slots is history.packed_slots assert score_view.meta.req_indices is history.local_req_indices assert score_view.meta.context_lens is history.context_lens diff --git a/tests/test_mla_attention_operator.py b/tests/test_mla_attention_operator.py index b7c2ea51..45885345 100644 --- a/tests/test_mla_attention_operator.py +++ b/tests/test_mla_attention_operator.py @@ -1,6 +1,6 @@ from __future__ import annotations -from unittest.mock import patch +from unittest.mock import Mock, patch import pytest import torch @@ -10,6 +10,7 @@ DecodeComputeView, ExplicitKVPayload, MlaLatentPayload, + PrefillComputeView, ) from sparsevllm.operators.mla_attention import ( MLA_ATTENTION_REGISTRY, @@ -265,6 +266,69 @@ def test_mla_provider_rejects_explicit_kv_before_kernel() -> None: kernel.assert_not_called() +def test_sgl_provider_uses_packed_varlen_prefill_metadata() -> None: + spec = _spec(tp_size=4) + workspace = _cpu_workspace(batch_size=1, head_count=5) + fa3 = Mock() + with ( + patch( + "sparsevllm.operators.mla_attention.allocate_mla_decode_workspace", + return_value=workspace, + ), + patch( + "sparsevllm.operators.mla_attention.SglFa3DecodeKernel", + return_value=fa3, + ), + ): + provider = MlaSglFa3Provider( + op_spec=spec, + device="cpu", + max_batch_size=1, + ) + q = torch.empty(2, 5, 256, dtype=torch.bfloat16) + output = torch.empty_like(q) + cu_seqlens_q = torch.tensor([0, 2], dtype=torch.int32) + cu_seqlens_k = torch.tensor([0, 4], dtype=torch.int32) + view = PrefillComputeView( + meta=AttentionViewMeta( + active_slots=torch.arange(4, dtype=torch.int32).view(1, 4), + req_indices=torch.tensor([0], dtype=torch.int32), + context_lens=torch.tensor([4], dtype=torch.int32), + max_context_len=4, + ), + payload=ExplicitKVPayload( + k_cache=torch.empty(4, 5, 256, dtype=torch.bfloat16), + v_cache=torch.empty(4, 5, 256, dtype=torch.bfloat16), + metadata={ + "layout": "mla_packed_varlen", + "cu_seqlens_k": cu_seqlens_k, + }, + ), + ) + fa3.run_contiguous_explicit_varlen.return_value = output + + actual = provider.run_explicit_prefill( + q, + view, + output, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=2, + ) + + assert actual is output + fa3.run_contiguous_explicit_varlen.assert_called_once_with( + q, + view.payload.k_cache, + view.payload.v_cache, + output, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=2, + max_seqlen_k=4, + ) + fa3.run_explicit_varlen.assert_not_called() + + def test_mla_provider_run_does_not_resolve_or_allocate() -> None: spec = _spec(tp_size=4) workspace = _cpu_workspace(batch_size=2, head_count=5) diff --git a/tests/test_moe_config.py b/tests/test_moe_config.py index 66dbc98b..868f59c2 100644 --- a/tests/test_moe_config.py +++ b/tests/test_moe_config.py @@ -404,3 +404,39 @@ def test_glm_h100_tp2_profile_covers_decode_and_large_prefill( config.block_k, config.group_m, ) == expected + + +@pytest.mark.parametrize( + ("num_tokens", "expected"), + [ + (1, (16, 128, 32, 8, 4, 4)), + (8, (16, 64, 128, 1, 4, 4)), + (128, (16, 64, 128, 1, 4, 4)), + (512, (128, 128, 64, 1, 8, 3)), + (65536, (128, 128, 64, 1, 8, 3)), + ], +) +def test_glm_h100_tp2_ep2_profile_covers_decode_and_long_prefill( + num_tokens, + expected, +): + common = dict( + dtype=torch.bfloat16, + num_tokens=num_tokens, + top_k=4, + num_local_experts=32, + hidden_size=2048, + intermediate_size=1536, + device_name="NVIDIA H100 80GB HBM3", + device_capability=(9, 0), + ) + for stage in ("w13", "w2"): + config = resolve_moe_gemm_config(**common, stage=stage) + assert ( + config.block_m, + config.block_n, + config.block_k, + config.group_m, + config.num_warps, + config.num_stages, + ) == expected diff --git a/tests/test_sgl_fa3.py b/tests/test_sgl_fa3.py index 9b109e5b..db576d95 100644 --- a/tests/test_sgl_fa3.py +++ b/tests/test_sgl_fa3.py @@ -290,7 +290,11 @@ def test_sgl_fa3_varlen_explicit_prefill_matches_causal_torch() -> None: k_cache = torch.randn( 2 * width, heads, head_dim, device=device, dtype=torch.bfloat16 ) - v_cache = torch.randn_like(k_cache) + v_backing = torch.randn( + 2 * width, heads, 448, device=device, dtype=torch.bfloat16 + ) + v_cache = v_backing[..., 192:] + assert not v_cache.is_contiguous() page_table = torch.arange( 2 * width, device=device, dtype=torch.int32 ).view(2, width) @@ -313,6 +317,25 @@ def test_sgl_fa3_varlen_explicit_prefill_matches_causal_torch() -> None: cu_seqlens_q=cu_seqlens_q, max_seqlen_q=max(chunk_lens), ) + packed_indices = torch.cat( + ( + page_table[1, : int(context_lens[0].item())], + page_table[0, : int(context_lens[1].item())], + ) + ).long() + packed_output = torch.empty_like(q) + kernel.run_contiguous_explicit_varlen( + q, + k_cache[packed_indices], + v_cache[packed_indices], + packed_output, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=torch.tensor( + [0, 5, 9], device=device, dtype=torch.int32 + ), + max_seqlen_q=max(chunk_lens), + max_seqlen_k=int(context_lens.max().item()), + ) expected_rows = [] query_start = 0 for batch_index, chunk_len in enumerate(chunk_lens): @@ -341,3 +364,4 @@ def test_sgl_fa3_varlen_explicit_prefill_matches_causal_torch() -> None: rtol=3e-2, atol=3e-2, ) + torch.testing.assert_close(packed_output, output) diff --git a/tests/test_tilelang_mla_operator.py b/tests/test_tilelang_mla_operator.py index e3865f91..1d8db2e7 100644 --- a/tests/test_tilelang_mla_operator.py +++ b/tests/test_tilelang_mla_operator.py @@ -338,6 +338,7 @@ def test_no_score_path_remains_fa3() -> None: tilelang.assert_not_called() fa3.assert_called_once() + assert fa3.call_args.kwargs["num_splits"] == 0 @pytest.mark.parametrize( From aa6e6a15e58cf123e3bc5462f72eacef867d9b01 Mon Sep 17 00:00:00 2001 From: CURRENTF <1206013760@qq.com> Date: Mon, 10 Aug 2026 04:50:29 +0800 Subject: [PATCH 20/48] fix: validate only real mla graph rows --- src/sparsevllm/layers/mla_attention.py | 9 +++- src/sparsevllm/operators/mla_attention.py | 14 ++++++ .../triton_kernel/mla/decode_schedule.py | 12 ++++- tests/test_mla_attention_layer.py | 49 +++++++++++++++++++ tests/test_mla_attention_operator.py | 1 + tests/test_mla_kernels.py | 32 ++++++++++++ 6 files changed, 115 insertions(+), 2 deletions(-) diff --git a/src/sparsevllm/layers/mla_attention.py b/src/sparsevllm/layers/mla_attention.py index eda27f45..04b14d5c 100644 --- a/src/sparsevllm/layers/mla_attention.py +++ b/src/sparsevllm/layers/mla_attention.py @@ -861,12 +861,19 @@ def run_decode( ) self._require_mla_payload(view, operation="MLA decode") output = torch.empty_like(q_nope_absorbed) + context = get_context() + valid_batch_size = ( + int(q_nope_absorbed.shape[0]) + if context.seqs is None + else len(context.seqs) + ) return self.provider.run( q_nope_absorbed, q_rope, view, output, - validation_scope=get_context().attention_validation_scope, + validation_scope=context.attention_validation_scope, + valid_batch_size=valid_batch_size, ) __all__ = [ diff --git a/src/sparsevllm/operators/mla_attention.py b/src/sparsevllm/operators/mla_attention.py index d1770e32..5694c6dd 100644 --- a/src/sparsevllm/operators/mla_attention.py +++ b/src/sparsevllm/operators/mla_attention.py @@ -95,6 +95,7 @@ def run( output: torch.Tensor, *, validation_scope: object | None = None, + valid_batch_size: int | None = None, ) -> torch.Tensor: raise NotImplementedError @@ -289,6 +290,7 @@ def _validate_metadata( payload: MlaLatentPayload, *, validation_scope: object | None, + valid_batch_size: int | None, ) -> None: cache_slot_count = int(payload.latent_cache.shape[0]) metadata_key = ( @@ -297,6 +299,7 @@ def _validate_metadata( view.meta.context_lens, cache_slot_count, view.meta.max_context_len, + valid_batch_size, ) cached = self._validated_decode_metadata cached_entries = ( @@ -312,6 +315,7 @@ def _validate_metadata( and entry[2] is metadata_key[2] and entry[3] == metadata_key[3] and entry[4] == metadata_key[4] + and entry[5] == metadata_key[5] for entry in cached_entries ) if metadata_is_validated: @@ -322,6 +326,7 @@ def _validate_metadata( view.meta.context_lens, cache_slot_count=cache_slot_count, max_context_len=view.meta.max_context_len, + valid_batch_size=valid_batch_size, ) if validation_scope is None: self._validated_decode_metadata = None @@ -358,6 +363,7 @@ def run( output: torch.Tensor, *, validation_scope: object | None = None, + valid_batch_size: int | None = None, ) -> torch.Tensor: payload = self._validate_run_inputs( q_nope_absorbed, @@ -369,6 +375,7 @@ def run( view, payload, validation_scope=validation_scope, + valid_batch_size=valid_batch_size, ) launch_config = self._launch_config_for( batch_size=int(q_nope_absorbed.shape[0]), @@ -442,6 +449,7 @@ def run( output: torch.Tensor, *, validation_scope: object | None = None, + valid_batch_size: int | None = None, ) -> torch.Tensor: if view.meta.attn_score is not None: return super().run( @@ -450,6 +458,7 @@ def run( view, output, validation_scope=validation_scope, + valid_batch_size=valid_batch_size, ) payload = self._validate_run_inputs( q_nope_absorbed, @@ -461,6 +470,7 @@ def run( view, payload, validation_scope=validation_scope, + valid_batch_size=valid_batch_size, ) return self.fa3( q_rope, @@ -695,6 +705,7 @@ def run( output: torch.Tensor, *, validation_scope: object | None = None, + valid_batch_size: int | None = None, ) -> torch.Tensor: attn_score = view.meta.attn_score if attn_score is None: @@ -704,6 +715,7 @@ def run( view, output, validation_scope=validation_scope, + valid_batch_size=valid_batch_size, ) # Per-head or non-tile-aligned score buffers remain on the existing # Triton implementation. This is a static shape dispatch before any @@ -724,6 +736,7 @@ def run( view, output, validation_scope=validation_scope, + valid_batch_size=valid_batch_size, ) payload = self._validate_run_inputs( q_nope_absorbed, @@ -735,6 +748,7 @@ def run( view, payload, validation_scope=validation_scope, + valid_batch_size=valid_batch_size, ) return self.tilelang_score( q_nope_absorbed, diff --git a/src/sparsevllm/triton_kernel/mla/decode_schedule.py b/src/sparsevllm/triton_kernel/mla/decode_schedule.py index 96086820..db838f01 100644 --- a/src/sparsevllm/triton_kernel/mla/decode_schedule.py +++ b/src/sparsevllm/triton_kernel/mla/decode_schedule.py @@ -319,6 +319,7 @@ def validate_mla_decode_metadata( *, cache_slot_count: int, max_context_len: int | None = None, + valid_batch_size: int | None = None, ) -> None: """Synchronously validate one decode view before per-layer reuse.""" @@ -343,6 +344,14 @@ def validate_mla_decode_metadata( if context_lens.ndim != 1 or context_lens.numel() == 0: raise ValueError("context_lens must be a non-empty one-dimensional tensor") batch_size = context_lens.numel() + if valid_batch_size is None: + valid_batch_size = batch_size + valid_batch_size = int(valid_batch_size) + if not 0 < valid_batch_size <= batch_size: + raise ValueError( + "valid_batch_size must be within the metadata batch: " + f"valid={valid_batch_size} batch={batch_size}" + ) if request_indices.shape != (batch_size,): raise ValueError( f"request_indices must have shape ({batch_size},), got " @@ -384,7 +393,8 @@ def validate_mla_decode_metadata( f"request_indices[{batch_index}]={request_row} is outside " f"[0, {active_slots.shape[0]})" ) - real_request_rows.append(request_row) + if batch_index < valid_batch_size: + real_request_rows.append(request_row) if length == 0: continue slots = active_slots[request_row, :length] diff --git a/tests/test_mla_attention_layer.py b/tests/test_mla_attention_layer.py index 873de62f..8729b161 100644 --- a/tests/test_mla_attention_layer.py +++ b/tests/test_mla_attention_layer.py @@ -9,6 +9,7 @@ from sparsevllm.engine.cache_manager import ( AttentionKeyComputeView, AttentionViewMeta, + DecodeComputeView, ExplicitKVPayload, MlaLatentPayload, PrefillComputeView, @@ -277,6 +278,54 @@ def test_set_context_starts_a_new_attention_validation_scope() -> None: assert second_step_scope is not first_step_scope +def test_mla_decode_passes_valid_batch_size_to_provider() -> None: + class _RecordingProvider(_TestProvider): + def run( + self, + q_nope_absorbed, + q_rope, + view, + output, + *, + validation_scope=None, + valid_batch_size=None, + ): + self.valid_batch_size = valid_batch_size + return output + + spec = _spec(tp_size=4) + provider = _RecordingProvider(spec, device="cpu", max_batch_size=4) + attention = MLAAttention( + spec=spec, + provider=provider, + prefill_workspace_bytes=64 * 1024 * 1024, + hidden_size=64, + projection_chunk_size=8, + ) + view = DecodeComputeView( + meta=AttentionViewMeta( + active_slots=torch.arange(16, dtype=torch.int32).view(4, 4), + req_indices=torch.tensor([0, 1, 2, 0], dtype=torch.int32), + context_lens=torch.tensor([4, 4, 4, 4], dtype=torch.int32), + ), + payload=MlaLatentPayload( + latent_cache=torch.empty(16, 1, 512, dtype=torch.bfloat16), + rope_cache=torch.empty(16, 1, 64, dtype=torch.bfloat16), + ), + ) + q_nope_absorbed = torch.empty(4, 5, 512, dtype=torch.bfloat16) + q_rope = torch.empty(4, 5, 64, dtype=torch.bfloat16) + + set_context(False, seqs=[object(), object(), object()]) + try: + output = attention.run_decode(q_nope_absorbed, q_rope, view) + finally: + reset_context() + + assert output.shape == q_nope_absorbed.shape + assert provider.valid_batch_size == 3 + + def test_mla_prefill_reuses_validated_packing_across_layers() -> None: attention = _attention() active_slots = torch.tensor([[3, 1]], dtype=torch.int32) diff --git a/tests/test_mla_attention_operator.py b/tests/test_mla_attention_operator.py index 45885345..56f86fa7 100644 --- a/tests/test_mla_attention_operator.py +++ b/tests/test_mla_attention_operator.py @@ -403,6 +403,7 @@ def test_mla_provider_run_does_not_resolve_or_allocate() -> None: view.meta.context_lens, cache_slot_count=4, max_context_len=None, + valid_batch_size=None, ) assert kernel.call_count == 3 kernel.assert_called_with( diff --git a/tests/test_mla_kernels.py b/tests/test_mla_kernels.py index 8f9de29b..6df057ce 100644 --- a/tests/test_mla_kernels.py +++ b/tests/test_mla_kernels.py @@ -19,6 +19,7 @@ prepare_mla_decode_schedule, run_mla_decode, select_glm_mla_decode_config, + validate_mla_decode_metadata, ) @@ -665,6 +666,37 @@ def test_mla_decode_rejects_duplicate_active_slots() -> None: ) +@CUDA_REQUIRED +def test_mla_decode_metadata_ignores_duplicate_graph_padding_rows() -> None: + active_slots = torch.tensor( + [[0, 1], [2, 3]], dtype=torch.int32, device="cuda" + ) + context_lens = torch.tensor([2, 2, 2, 2], dtype=torch.int32, device="cuda") + padded_request_indices = torch.tensor( + [0, 1, 0, 0], dtype=torch.int32, device="cuda" + ) + + validate_mla_decode_metadata( + active_slots, + padded_request_indices, + context_lens, + cache_slot_count=4, + valid_batch_size=2, + ) + + duplicate_real_request_indices = torch.tensor( + [0, 0, 0, 0], dtype=torch.int32, device="cuda" + ) + with pytest.raises(ValueError, match="duplicate non-padding rows"): + validate_mla_decode_metadata( + active_slots, + duplicate_real_request_indices, + context_lens, + cache_slot_count=4, + valid_batch_size=2, + ) + + @CUDA_REQUIRED def test_decode_stage1_matches_per_block_oracle() -> None: torch.manual_seed(23) From 8cfadba41056c23cb0c6747ff4f41d1946cca795 Mon Sep 17 00:00:00 2001 From: CURRENTF <1206013760@qq.com> Date: Mon, 10 Aug 2026 13:20:05 +0800 Subject: [PATCH 21/48] perf: optimize glm tp2 ep2 decode --- .../flashinfer_custom_all_reduce.py | 76 ++++++-- .../distributed/parallel_context.py | 25 ++- src/sparsevllm/engine/model_runner.py | 2 - src/sparsevllm/layers/activation.py | 17 +- src/sparsevllm/models/glm4_moe_lite.py | 13 +- src/sparsevllm/models/qwen3.py | 4 +- src/sparsevllm/operators/activation.py | 147 +++++++++++++++ src/sparsevllm/operators/moe.py | 20 +- src/sparsevllm/operators/sgl_moe.py | 65 +++++-- src/sparsevllm/triton_kernel/moe.py | 60 ++++++ tests/test_activation.py | 67 +++++++ tests/test_operator_providers.py | 99 ++++++++++ tests/test_parallel_context.py | 93 +++++++++- tests/test_sgl_moe.py | 172 ++++++++++++++++++ 14 files changed, 806 insertions(+), 54 deletions(-) create mode 100644 src/sparsevllm/operators/activation.py create mode 100644 tests/test_activation.py diff --git a/src/sparsevllm/distributed/flashinfer_custom_all_reduce.py b/src/sparsevllm/distributed/flashinfer_custom_all_reduce.py index 9bc8761c..d52cf5bf 100644 --- a/src/sparsevllm/distributed/flashinfer_custom_all_reduce.py +++ b/src/sparsevllm/distributed/flashinfer_custom_all_reduce.py @@ -16,6 +16,65 @@ class FlashInferCustomAllReduce: """Small-tensor TP all-reduce using FlashInfer's vLLM kernel.""" + _REQUIRED_COMM_APIS = ( + "CudaRTLibrary", + "create_shared_buffer", + "vllm_all_reduce", + "vllm_dispose", + "vllm_get_graph_buffer_ipc_meta", + "vllm_init_custom_ar", + "vllm_meta_size", + "vllm_register_buffer", + "vllm_register_graph_buffers", + ) + + @classmethod + def unsupported_reason(cls, group: ParallelGroup) -> str | None: + """Return why this group cannot use the provider, or ``None``.""" + + if group.size != 2 or group.process_group is None: + return "requires a two-rank process group" + if dist.get_backend(group.process_group) != dist.Backend.NCCL: + return "requires an NCCL process group" + if not torch.cuda.is_available(): + return "requires CUDA" + + device_count = int(torch.cuda.device_count()) + if any(rank < 0 or rank >= device_count for rank in group.ranks): + return ( + "requires all group ranks to map to visible local CUDA devices, " + f"got ranks={group.ranks}, device_count={device_count}" + ) + current_device = int(torch.cuda.current_device()) + capability = tuple( + int(value) for value in torch.cuda.get_device_capability(current_device) + ) + if capability != (9, 0): + return f"requires CUDA SM90, got {capability}" + expected_device = int(group.ranks[group.rank]) + if current_device != expected_device: + return ( + "requires world-rank-to-device mapping for CUDA IPC, " + f"got current_device={current_device}, expected={expected_device}" + ) + for peer_rank, peer_device in enumerate(group.ranks): + if peer_rank == group.rank: + continue + if not torch.cuda.can_device_access_peer(current_device, peer_device): + return ( + "requires CUDA peer access between all ranks, " + f"missing {current_device}->{peer_device} for ranks={group.ranks}" + ) + + try: + from flashinfer import comm + except (ImportError, OSError, RuntimeError) as exc: + return f"FlashInfer communication APIs are unavailable: {exc}" + missing = [name for name in cls._REQUIRED_COMM_APIS if not hasattr(comm, name)] + if missing: + return "FlashInfer communication APIs are missing: " + ", ".join(missing) + return None + def __init__( self, group: ParallelGroup, @@ -23,25 +82,16 @@ def __init__( max_size_bytes: int = 8 * 1024 * 1024, num_ctas: int = 32, ) -> None: - if group.size != 2 or group.process_group is None: - raise ValueError( - "FlashInfer custom all-reduce requires a two-rank process group." - ) - if dist.get_backend(group.process_group) != dist.Backend.NCCL: - raise ValueError("FlashInfer custom all-reduce requires an NCCL group.") if max_size_bytes <= 0 or num_ctas <= 0: raise ValueError( "FlashInfer custom all-reduce sizes must be positive, got " f"max_size_bytes={max_size_bytes}, num_ctas={num_ctas}." ) - if not all( - torch.cuda.can_device_access_peer(group.rank, peer_rank) - for peer_rank in range(group.size) - if peer_rank != group.rank - ): + unsupported_reason = self.unsupported_reason(group) + if unsupported_reason is not None: raise RuntimeError( - "FlashInfer custom all-reduce requires CUDA peer access between " - f"all ranks in {group.ranks}." + "FlashInfer custom all-reduce is unsupported: " + f"{unsupported_reason}." ) from flashinfer.comm import ( diff --git a/src/sparsevllm/distributed/parallel_context.py b/src/sparsevllm/distributed/parallel_context.py index acacadca..43a4c6b2 100644 --- a/src/sparsevllm/distributed/parallel_context.py +++ b/src/sparsevllm/distributed/parallel_context.py @@ -9,6 +9,7 @@ from sparsevllm.distributed.topology import ParallelTopology, parallel_group_ranks from sparsevllm.operators.all_reduce import AllReduceProvider, resolve_all_reduce_provider +from sparsevllm.utils.log import logger class GraphAllReduceProvider(Protocol): @@ -330,10 +331,26 @@ def bind_provider(group: ParallelGroup) -> ParallelGroup: FlashInferCustomAllReduce, ) - _PARALLEL_CONTEXT = replace( - _PARALLEL_CONTEXT, - all_reduce_provider=FlashInferCustomAllReduce(_PARALLEL_CONTEXT.tensor), - ) + group = _PARALLEL_CONTEXT.tensor + local_reason = FlashInferCustomAllReduce.unsupported_reason(group) + rank_reasons = [None] * group.size + dist.all_gather_object(rank_reasons, local_reason, group=group.process_group) + unsupported = [ + (rank, reason) + for rank, reason in zip(group.ranks, rank_reasons) + if reason is not None + ] + if unsupported: + logger.info( + "Using NCCL TP all-reduce because FlashInfer custom all-reduce " + "is unsupported on ranks {}.", + unsupported, + ) + else: + _PARALLEL_CONTEXT = replace( + _PARALLEL_CONTEXT, + all_reduce_provider=FlashInferCustomAllReduce(group), + ) return _PARALLEL_CONTEXT diff --git a/src/sparsevllm/engine/model_runner.py b/src/sparsevllm/engine/model_runner.py index 23a3f70e..565340e8 100644 --- a/src/sparsevllm/engine/model_runner.py +++ b/src/sparsevllm/engine/model_runner.py @@ -153,10 +153,8 @@ def __init__( hf_config.model_type == "glm4_moe_lite" and bool(config.decode_cuda_graph) and config.tensor_parallel_size == 2 - and config.expert_parallel_size == 1 and config.data_parallel_size == 1 and self.platform.is_cuda() - and self.platform.get_device_caps(rank).compute_capability == (9, 0) ), ) if self.parallel_context.all_reduce_provider is not None: diff --git a/src/sparsevllm/layers/activation.py b/src/sparsevllm/layers/activation.py index 2e06e5ca..393c609f 100755 --- a/src/sparsevllm/layers/activation.py +++ b/src/sparsevllm/layers/activation.py @@ -1,18 +1,17 @@ import torch from torch import nn -import torch.nn.functional as F + +from sparsevllm.operators.activation import ( + SiluAndMulProvider, + TorchSiluAndMulProvider, +) class SiluAndMul(nn.Module): - def __init__(self): + def __init__(self, provider: SiluAndMulProvider | None = None): super().__init__() + self.provider = provider if provider is not None else TorchSiluAndMulProvider() def forward(self, x: torch.Tensor) -> torch.Tensor: - # In-place activation to reduce peak memory: - # x is typically (num_tokens, 2 * intermediate_size) in prefill; allocating an extra - # (num_tokens, intermediate_size) output can be multiple GiB at long-context, large-batch. - x, y = x.chunk(2, -1) - F.silu(x, inplace=True) - x.mul_(y) - return x + return self.provider(x) diff --git a/src/sparsevllm/models/glm4_moe_lite.py b/src/sparsevllm/models/glm4_moe_lite.py index 0a4fa191..50b4b7a6 100644 --- a/src/sparsevllm/models/glm4_moe_lite.py +++ b/src/sparsevllm/models/glm4_moe_lite.py @@ -30,6 +30,7 @@ from sparsevllm.layers.rotary_embedding import RotaryEmbedding, get_rope from sparsevllm.models.qwen3 import Qwen3MLP from sparsevllm.operators.mla_attention import MlaAttentionOpSpec +from sparsevllm.operators.activation import resolve_silu_and_mul_provider from sparsevllm.operators.moe import ( MoeOpSpec, model_activation_dtype, @@ -472,7 +473,11 @@ def __init__( self.shared_expert_id = ( self.routed_num_experts if self.fuses_shared_decode else None ) - self.shared_act = SiluAndMul() + self.shared_act = SiluAndMul( + provider=resolve_silu_and_mul_provider( + activation_dtype=model_activation_dtype(config), + ) + ) if self.fuses_shared_decode: self.routed_op_spec = MoeOpSpec( num_experts=self.routed_num_experts, @@ -578,6 +583,9 @@ def __init__( mlp_chunk_size=self.mlp_chunk_size, quantization=None, reduce_results=False, + activation_provider=resolve_silu_and_mul_provider( + activation_dtype=model_activation_dtype(config), + ), ) ) @@ -763,6 +771,9 @@ def __init__( hidden_act=str(config.hidden_act), mlp_chunk_size=int(mlp_chunk_size), quantization=None, + activation_provider=resolve_silu_and_mul_provider( + activation_dtype=model_activation_dtype(config), + ), ) elif layer_type == "sparse": self.mlp = Glm4MoeLiteSparseMoeBlock( diff --git a/src/sparsevllm/models/qwen3.py b/src/sparsevllm/models/qwen3.py index 898d3116..071a8ee3 100755 --- a/src/sparsevllm/models/qwen3.py +++ b/src/sparsevllm/models/qwen3.py @@ -7,6 +7,7 @@ from sparsevllm.utils.context import get_context from sparsevllm.layers.activation import SiluAndMul +from sparsevllm.operators.activation import SiluAndMulProvider from sparsevllm.layers.attention import Attention from sparsevllm.layers.layernorm import RMSNorm from sparsevllm.layers.linear import QKVParallelLinear, MergedColumnParallelLinear, RowParallelLinear @@ -170,6 +171,7 @@ def __init__( mlp_chunk_size: int = 16384, quantization=None, reduce_results: bool = True, + activation_provider: SiluAndMulProvider | None = None, ) -> None: super().__init__() self.gate_up_proj = MergedColumnParallelLinear( @@ -186,7 +188,7 @@ def __init__( reduce_results=reduce_results, ) assert hidden_act == "silu" - self.act_fn = SiluAndMul() + self.act_fn = SiluAndMul(provider=activation_provider) self.mlp_chunk_size = int(mlp_chunk_size) if self.mlp_chunk_size <= 0: raise ValueError(f"mlp_chunk_size must be > 0, got {mlp_chunk_size}.") diff --git a/src/sparsevllm/operators/activation.py b/src/sparsevllm/operators/activation.py new file mode 100644 index 00000000..dea72f33 --- /dev/null +++ b/src/sparsevllm/operators/activation.py @@ -0,0 +1,147 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import torch +import torch.nn.functional as F + +import sparsevllm.platforms as platforms +from sparsevllm.operators.registry import OpRegistry, OpResolver, SupportResult +from sparsevllm.platforms.interface import DeviceCaps, PlatformEnum + + +@dataclass(frozen=True) +class SiluAndMulSpec: + activation_dtype: torch.dtype + input_ndim: int = 2 + contiguous: bool = True + + def __post_init__(self) -> None: + if int(self.input_ndim) <= 0: + raise ValueError("SiluAndMul input_ndim must be positive.") + + +def _validate_input(x: torch.Tensor) -> None: + if int(x.shape[-1]) % 2: + raise ValueError( + "SiluAndMul requires an even final dimension, got " + f"{int(x.shape[-1])}." + ) + + +def _validate_bound_input(x: torch.Tensor, spec: SiluAndMulSpec) -> None: + _validate_input(x) + if x.dtype != spec.activation_dtype: + raise TypeError( + "Bound SiluAndMul provider requires " + f"dtype={spec.activation_dtype}, got {x.dtype}." + ) + if x.ndim != int(spec.input_ndim): + raise ValueError( + "Bound SiluAndMul provider requires " + f"ndim={spec.input_ndim}, got {x.ndim}." + ) + if spec.contiguous and not x.is_contiguous(): + raise ValueError("Bound SiluAndMul provider requires contiguous input.") + + +class SiluAndMulProvider: + name = "" + priority = 0 + + def __call__(self, x: torch.Tensor) -> torch.Tensor: + raise NotImplementedError + + +SILU_AND_MUL_REGISTRY: OpRegistry[SiluAndMulSpec, SiluAndMulProvider] = OpRegistry( + "SiLU-and-multiply" +) + + +@SILU_AND_MUL_REGISTRY.register +class TritonSiluAndMulProvider(SiluAndMulProvider): + name = "triton" + priority = 10 + + def __init__(self, *, op_spec: SiluAndMulSpec) -> None: + self.spec = op_spec + + @classmethod + def supports(cls, spec: SiluAndMulSpec, caps: DeviceCaps) -> SupportResult: + if caps.platform != PlatformEnum.CUDA: + return SupportResult.no(f"requires CUDA, got {caps.platform.name}") + if not caps.supports_triton: + return SupportResult.no("platform does not support Triton") + if spec.activation_dtype not in (torch.float16, torch.bfloat16): + return SupportResult.no( + "requires FP16 or BF16 activations, " + f"got {spec.activation_dtype}" + ) + if int(spec.input_ndim) != 2 or not spec.contiguous: + return SupportResult.no("requires contiguous rank-2 inputs") + return SupportResult.yes() + + def __call__(self, x: torch.Tensor) -> torch.Tensor: + _validate_bound_input(x, self.spec) + if not x.is_cuda: + raise ValueError("Triton SiluAndMul provider requires a CUDA input.") + from sparsevllm.triton_kernel.silu_and_mul import silu_and_mul_fwd + + return silu_and_mul_fwd(x) + + +@SILU_AND_MUL_REGISTRY.register +class TorchSiluAndMulProvider(SiluAndMulProvider): + name = "torch" + priority = 0 + + def __init__(self, *, op_spec: SiluAndMulSpec | None = None) -> None: + self.spec = op_spec + + @classmethod + def supports(cls, spec: SiluAndMulSpec, caps: DeviceCaps) -> SupportResult: + del spec, caps + return SupportResult.yes() + + def __call__(self, x: torch.Tensor) -> torch.Tensor: + if self.spec is None: + _validate_input(x) + else: + _validate_bound_input(x, self.spec) + gate, up = x.chunk(2, -1) + F.silu(gate, inplace=True) + gate.mul_(up) + return gate + + +def resolve_silu_and_mul_provider( + *, + activation_dtype: torch.dtype, + input_ndim: int = 2, + contiguous: bool = True, + device_index: int | None = None, +) -> SiluAndMulProvider: + platform = platforms.current_platform + if device_index is None: + device_index = torch.cuda.current_device() if platform.is_cuda_alike() else 0 + caps = platform.get_device_caps(int(device_index)) + spec = SiluAndMulSpec( + activation_dtype=activation_dtype, + input_ndim=int(input_ndim), + contiguous=bool(contiguous), + ) + return OpResolver(SILU_AND_MUL_REGISTRY).resolve( + spec, + caps, + op_spec=spec, + ).provider + + +__all__ = [ + "SILU_AND_MUL_REGISTRY", + "SiluAndMulProvider", + "SiluAndMulSpec", + "TorchSiluAndMulProvider", + "TritonSiluAndMulProvider", + "resolve_silu_and_mul_provider", +] diff --git a/src/sparsevllm/operators/moe.py b/src/sparsevllm/operators/moe.py index 5267675a..3ee15e7b 100644 --- a/src/sparsevllm/operators/moe.py +++ b/src/sparsevllm/operators/moe.py @@ -191,6 +191,7 @@ def supports(cls, spec: MoeOpSpec, caps: DeviceCaps) -> SupportResult: expected = { (64, 64, 2048, 768, 4, 2, 1, "biased_sigmoid"), (65, 65, 2048, 768, 5, 2, 1, "biased_sigmoid"), + (64, 32, 2048, 1536, 4, 1, 2, "biased_sigmoid"), } actual = ( spec.num_experts, @@ -210,9 +211,15 @@ def supports(cls, spec: MoeOpSpec, caps: DeviceCaps) -> SupportResult: return SupportResult.no("requires BF16 activations") if spec.weight_dtype != torch.bfloat16 or spec.block_shape is not None: return SupportResult.no("requires unquantized BF16 expert weights") - from sparsevllm.operators.sgl_moe import sgl_moe_alignment_support + from sparsevllm.operators.sgl_moe import ( + sgl_moe_alignment_support, + sgl_moe_ep_alignment_support, + ) - supported, reason = sgl_moe_alignment_support() + if spec.ep_size > 1: + supported, reason = sgl_moe_ep_alignment_support() + else: + supported, reason = sgl_moe_alignment_support() return SupportResult.yes() if supported else SupportResult.no(reason) def run( @@ -235,6 +242,13 @@ def run( from sparsevllm.operators.sgl_moe import sgl_moe_align_block_size from sparsevllm.triton_kernel.moe import fused_moe + alignment_impl = None + num_tokens = int(hidden_states.shape[0]) + if ( + num_tokens <= 64 + and num_tokens * int(spec.top_k) * 4 > int(spec.num_experts) + ): + alignment_impl = sgl_moe_align_block_size return fused_moe( hidden_states, w13_weight, @@ -243,7 +257,7 @@ def run( topk_weights, num_experts=spec.num_experts, local_expert_start=local_expert_start, - alignment_impl=sgl_moe_align_block_size, + alignment_impl=alignment_impl, ) diff --git a/src/sparsevllm/operators/sgl_moe.py b/src/sparsevllm/operators/sgl_moe.py index 5842d6a1..d574cbbb 100644 --- a/src/sparsevllm/operators/sgl_moe.py +++ b/src/sparsevllm/operators/sgl_moe.py @@ -80,6 +80,12 @@ def sgl_moe_alignment_support() -> tuple[bool, str]: return _sgl_moe_support("MoE alignment") +def sgl_moe_ep_alignment_support() -> tuple[bool, str]: + """Check the invalid-route filtering required by expert parallelism.""" + + return _sgl_moe_support("MoE EP alignment") + + class SglGlmFusedMoeGate: """One-kernel sigmoid, biased top-k, normalization, and route scaling.""" @@ -168,13 +174,42 @@ def sgl_moe_align_block_size( local_expert_start: int, local_expert_end: int, ) -> MoeAlignment: - """Build a full-expert MoE assignment with the SGL CUDA kernel.""" + """Build a local MoE assignment with the SGL CUDA kernel.""" + + local_expert_start = int(local_expert_start) + local_expert_end = int(local_expert_end) + num_local_experts = local_expert_end - local_expert_start + if not 0 <= local_expert_start < local_expert_end <= int(num_experts): + raise ValueError( + "Invalid local expert range " + f"[{local_expert_start}, {local_expert_end}) for {num_experts} experts." + ) + supported, reason = sgl_moe_alignment_support() + if not supported: + raise RuntimeError(reason) + installed = _installed_sgl_kernel_version() + if installed is None: + raise RuntimeError("sgl-kernel package metadata is unavailable") + _, parsed = installed + + local_topk_ids = topk_ids + has_remote_experts = num_local_experts != int(num_experts) + ignore_invalid_expert = has_remote_experts and parsed[:2] == (0, 4) + if num_local_experts != int(num_experts): + from sparsevllm.triton_kernel.moe import localize_expert_ids - if (int(local_expert_start), int(local_expert_end)) != (0, int(num_experts)): - raise ValueError("SGL MoE alignment currently requires EP size 1.") - num_assignments = int(topk_ids.numel()) + local_topk_ids = localize_expert_ids( + topk_ids, + local_expert_start=local_expert_start, + local_expert_end=local_expert_end, + remote_expert_id=-1 if ignore_invalid_expert else num_local_experts, + ) + num_assignments = int(local_topk_ids.numel()) + padding_experts = num_local_experts + int( + has_remote_experts and parsed[:2] == (0, 3) + ) max_num_tokens_padded = triton.cdiv( - num_assignments + int(num_experts) * (int(block_size) - 1), + num_assignments + padding_experts * (int(block_size) - 1), int(block_size), ) * int(block_size) sorted_token_ids = torch.empty( @@ -193,25 +228,18 @@ def sgl_moe_align_block_size( device=topk_ids.device, ) cumsum_buffer = torch.empty( - int(num_experts) + 1, + num_local_experts + 1, dtype=torch.int32, device=topk_ids.device, ) - supported, reason = sgl_moe_alignment_support() - if not supported: - raise RuntimeError(reason) - installed = _installed_sgl_kernel_version() - if installed is None: - raise RuntimeError("sgl-kernel package metadata is unavailable") - _, parsed = installed from sgl_kernel import moe_align_block_size # Both validated APIs iterate to num_experts - 1. The extra empty logical # expert makes the complete [0, num_experts) range participate. if parsed[:2] == (0, 3): moe_align_block_size( - topk_ids, - int(num_experts) + 1, + local_topk_ids, + num_local_experts + 1, int(block_size), sorted_token_ids, expert_ids, @@ -221,15 +249,15 @@ def sgl_moe_align_block_size( ) else: moe_align_block_size( - topk_ids, - int(num_experts) + 1, + local_topk_ids, + num_local_experts + 1, int(block_size), sorted_token_ids, expert_ids, num_tokens_post_padded, cumsum_buffer, True, - False, + ignore_invalid_expert, ) return MoeAlignment( sorted_token_ids=sorted_token_ids, @@ -244,5 +272,6 @@ def sgl_moe_align_block_size( "SglGlmFusedMoeGate", "sgl_fused_moe_gate_support", "sgl_moe_alignment_support", + "sgl_moe_ep_alignment_support", "sgl_moe_align_block_size", ] diff --git a/src/sparsevllm/triton_kernel/moe.py b/src/sparsevllm/triton_kernel/moe.py index a2b93303..7b83cd9f 100644 --- a/src/sparsevllm/triton_kernel/moe.py +++ b/src/sparsevllm/triton_kernel/moe.py @@ -28,6 +28,66 @@ class MoeAlignment: naive: bool +@triton.jit +def _localize_expert_ids_kernel( + input_ids_ptr, + output_ids_ptr, + num_assignments, + local_expert_start: tl.constexpr, + local_expert_end: tl.constexpr, + remote_expert_id: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + valid = offsets < num_assignments + expert_ids = tl.load(input_ids_ptr + offsets, mask=valid, other=-1) + local_ids = tl.where( + (expert_ids >= local_expert_start) & (expert_ids < local_expert_end), + expert_ids - local_expert_start, + remote_expert_id, + ) + tl.store(output_ids_ptr + offsets, local_ids, mask=valid) + + +def localize_expert_ids( + expert_ids: torch.Tensor, + *, + local_expert_start: int, + local_expert_end: int, + remote_expert_id: int = -1, +) -> torch.Tensor: + """Map global expert IDs to one EP shard and encode remote routes.""" + + if not expert_ids.is_cuda or not expert_ids.is_contiguous(): + raise ValueError("Expert ID localization requires a contiguous CUDA tensor.") + if expert_ids.dtype not in (torch.int32, torch.int64): + raise TypeError("Expert ID localization requires int32 or int64 IDs.") + if expert_ids.numel() <= 0: + raise ValueError("Expert ID localization requires at least one ID.") + local_expert_start = int(local_expert_start) + local_expert_end = int(local_expert_end) + remote_expert_id = int(remote_expert_id) + if not 0 <= local_expert_start < local_expert_end: + raise ValueError( + "Invalid local expert range " + f"[{local_expert_start}, {local_expert_end})." + ) + output = torch.empty_like(expert_ids, dtype=torch.int32) + block_size = 256 + _localize_expert_ids_kernel[ + (triton.cdiv(int(expert_ids.numel()), block_size),) + ]( + expert_ids, + output, + int(expert_ids.numel()), + local_expert_start=local_expert_start, + local_expert_end=local_expert_end, + remote_expert_id=remote_expert_id, + BLOCK_SIZE=block_size, + ) + return output + + @triton.jit def _append_shared_expert_route_kernel( input_ids_ptr, diff --git a/tests/test_activation.py b/tests/test_activation.py new file mode 100644 index 00000000..bfa6b055 --- /dev/null +++ b/tests/test_activation.py @@ -0,0 +1,67 @@ +import pytest +import torch +import torch.nn.functional as F + +from sparsevllm.layers.activation import SiluAndMul +from sparsevllm.operators.activation import ( + SiluAndMulSpec, + TorchSiluAndMulProvider, + TritonSiluAndMulProvider, +) + + +def _reference(x: torch.Tensor) -> torch.Tensor: + gate, up = x.chunk(2, dim=-1) + return F.silu(gate) * up + + +def test_silu_and_mul_cpu_matches_reference(): + x = torch.randn(3, 16, dtype=torch.float32) + + torch.testing.assert_close( + SiluAndMul(provider=TorchSiluAndMulProvider())(x.clone()), + _reference(x), + ) + + +def test_silu_and_mul_rejects_odd_width(): + with pytest.raises(ValueError, match="even final dimension"): + SiluAndMul(provider=TorchSiluAndMulProvider())(torch.randn(2, 7)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("rows", [8, 256, 257]) +def test_silu_and_mul_cuda_matches_reference_and_aliases_input( + dtype: torch.dtype, + rows: int, +): + torch.manual_seed(20260810) + x = torch.randn(rows, 3072, device="cuda", dtype=dtype) + expected = _reference(x) + actual_input = x.clone() + up_before = actual_input[:, 1536:].clone() + actual = SiluAndMul( + provider=TritonSiluAndMulProvider( + op_spec=SiluAndMulSpec(activation_dtype=dtype), + ) + )(actual_input) + + torch.testing.assert_close(actual, expected, rtol=2e-3, atol=2e-3) + assert actual.data_ptr() == actual_input.data_ptr() + torch.testing.assert_close(actual_input[:, :1536], actual) + torch.testing.assert_close(actual_input[:, 1536:], up_before) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_bound_triton_silu_and_mul_rejects_contract_mismatch(): + provider = TritonSiluAndMulProvider( + op_spec=SiluAndMulSpec(activation_dtype=torch.bfloat16), + ) + + with pytest.raises(TypeError, match="requires dtype"): + provider(torch.randn(4, 16, dtype=torch.float32, device="cuda")) + with pytest.raises(ValueError, match="contiguous"): + provider( + torch.randn(4, 16, dtype=torch.bfloat16, device="cuda").transpose(0, 1) + ) diff --git a/tests/test_operator_providers.py b/tests/test_operator_providers.py index 0748300d..7b4a4537 100644 --- a/tests/test_operator_providers.py +++ b/tests/test_operator_providers.py @@ -6,6 +6,12 @@ import torch from sparsevllm.operators.all_reduce import HopperTp2FlashInferAllReduceProvider +from sparsevllm.operators.activation import ( + SILU_AND_MUL_REGISTRY, + SiluAndMulSpec, + TorchSiluAndMulProvider, + TritonSiluAndMulProvider, +) from sparsevllm.operators.fp8_linear import ( FP8_LINEAR_REGISTRY, FlashInferSm90Fp8LinearProvider, @@ -23,6 +29,7 @@ FlashInferCutlassFp8MoeProvider, HopperQwen36HybridFp8MoeProvider, MoeOpSpec, + SglAlignedTritonGlmMoeProvider, resolve_moe_provider, ) from sparsevllm.operators.registry import OpResolver @@ -192,6 +199,22 @@ def test_native_gate_up_provider_matches_swiglu_semantics(): torch.testing.assert_close(actual, expected) +def test_silu_and_mul_provider_respects_platform_capability(): + cuda_provider = OpResolver(SILU_AND_MUL_REGISTRY).resolve( + SiluAndMulSpec(activation_dtype=torch.bfloat16), + _cuda_caps((9, 0)), + op_spec=SiluAndMulSpec(activation_dtype=torch.bfloat16), + ).provider + cpu_provider = OpResolver(SILU_AND_MUL_REGISTRY).resolve( + SiluAndMulSpec(activation_dtype=torch.float32), + _non_cuda_caps(PlatformEnum.CPU), + op_spec=SiluAndMulSpec(activation_dtype=torch.float32), + ).provider + + assert isinstance(cuda_provider, TritonSiluAndMulProvider) + assert isinstance(cpu_provider, TorchSiluAndMulProvider) + + @pytest.mark.parametrize( "overrides", [ @@ -515,6 +538,82 @@ def test_glm_tp2_fused_shared_decode_uses_sgl_aligned_provider(): assert resolved.provider.name == "sgl_aligned_triton_glm" +def test_glm_tp2_ep2_moe_uses_sgl_aligned_provider(): + spec = _moe_spec( + activation_dtype=torch.bfloat16, + weight_dtype=torch.bfloat16, + block_shape=None, + hidden_size=2048, + intermediate_size=1536, + num_local_experts=32, + num_experts=64, + top_k=4, + ep_size=2, + tp_size=1, + routing_method="biased_sigmoid", + ) + + with patch( + "sparsevllm.operators.sgl_moe.sgl_moe_ep_alignment_support", + return_value=(True, "supported"), + ): + resolved = OpResolver(MOE_REGISTRY).resolve( + spec, + _cuda_caps( + (9, 0), + native_fp8=False, + device_name="NVIDIA H100 80GB HBM3", + ), + ) + + assert resolved.provider.name == "sgl_aligned_triton_glm" + + +@pytest.mark.parametrize( + ("num_tokens", "expects_sgl_alignment"), + [(4, False), (5, True), (64, True), (65, False)], +) +def test_glm_tp2_ep2_sgl_alignment_is_bounded( + num_tokens, + expects_sgl_alignment, +): + spec = _moe_spec( + activation_dtype=torch.bfloat16, + weight_dtype=torch.bfloat16, + block_shape=None, + hidden_size=2048, + intermediate_size=1536, + num_local_experts=32, + num_experts=64, + top_k=4, + ep_size=2, + tp_size=1, + routing_method="biased_sigmoid", + ) + hidden_states = torch.empty(num_tokens, 2048) + topk_ids = torch.empty(num_tokens, 4, dtype=torch.int64) + topk_weights = torch.empty(num_tokens, 4) + weights = torch.empty(0) + provider = SglAlignedTritonGlmMoeProvider() + + with patch("sparsevllm.triton_kernel.moe.fused_moe") as fused_moe: + provider.run( + spec, + hidden_states, + topk_ids, + topk_weights, + weights, + weights, + None, + None, + local_expert_start=0, + ep_rank=0, + ) + + alignment_impl = fused_moe.call_args.kwargs["alignment_impl"] + assert (alignment_impl is not None) is expects_sgl_alignment + + @pytest.mark.parametrize( ("tp_size", "ep_size", "intermediate_size", "num_local_experts"), [(1, 1, 512, 256), (2, 1, 256, 256), (1, 2, 512, 128)], diff --git a/tests/test_parallel_context.py b/tests/test_parallel_context.py index ae59b857..bd256c61 100644 --- a/tests/test_parallel_context.py +++ b/tests/test_parallel_context.py @@ -12,13 +12,18 @@ ParallelGroup, ParallelMode, ParallelTopology, - get_parallel_context, - init_parallel_context, parallel_group_ranks, parallel_ranks_from_world_rank, - reset_parallel_context, world_rank_from_parallel_ranks, ) +from sparsevllm.distributed.flashinfer_custom_all_reduce import ( + FlashInferCustomAllReduce, +) +from sparsevllm.distributed.parallel_context import ( + get_parallel_context, + init_parallel_context, + reset_parallel_context, +) from sparsevllm.engine.cache_manager.base import CacheManager from sparsevllm.layers.embed_head import VocabParallelEmbedding from sparsevllm.layers.linear import ColumnParallelLinear, RowParallelLinear @@ -158,6 +163,88 @@ def test_hybrid_moe_parallel_context_uses_explicit_groups(): reset_parallel_context() +@pytest.mark.parametrize( + ("local_reason", "rank_reasons", "expects_provider"), + [ + (None, [None, None], True), + ("peer access is unavailable", ["peer access is unavailable"] * 2, False), + (None, [None, "peer access is unavailable"], False), + ], +) +def test_hybrid_tp2_ep2_resolves_custom_all_reduce_at_init( + local_reason, + rank_reasons, + expects_provider, +): + reset_parallel_context() + provider = Mock() + provider_cls = Mock(return_value=provider) + provider_cls.unsupported_reason.return_value = local_reason + + def all_gather_object(output, _local_reason, *, group): + assert group is dist.group.WORLD + output[:] = rank_reasons + + with ( + patch.object(dist, "is_initialized", return_value=True), + patch.object(dist, "get_world_size", return_value=2), + patch.object(dist, "get_rank", return_value=0), + patch.object(dist, "all_gather_object", side_effect=all_gather_object), + patch( + "sparsevllm.distributed.flashinfer_custom_all_reduce." + "FlashInferCustomAllReduce", + provider_cls, + ), + ): + context = init_parallel_context( + topology=ParallelTopology(2, 2, 1, ParallelMode.OUTER_TP_MOE), + enable_flashinfer_custom_all_reduce=True, + ) + + selected_group = provider_cls.unsupported_reason.call_args.args[0] + assert selected_group.ranks == (0, 1) + assert context.attention.ranks == (0, 1) + assert context.expert.ranks == (0, 1) + assert context.moe_tensor.ranks == (0,) + if expects_provider: + provider_cls.assert_called_once_with(selected_group) + assert context.all_reduce_provider is provider + else: + provider_cls.assert_not_called() + assert context.all_reduce_provider is None + reset_parallel_context() + + +def test_custom_all_reduce_support_rejects_missing_peer_access(): + group = ParallelGroup(object(), (0, 1), 0, 2) + with ( + patch.object(dist, "get_backend", return_value=dist.Backend.NCCL), + patch.object(torch.cuda, "is_available", return_value=True), + patch.object(torch.cuda, "device_count", return_value=2), + patch.object(torch.cuda, "current_device", return_value=0), + patch.object(torch.cuda, "get_device_capability", return_value=(9, 0)), + patch.object(torch.cuda, "can_device_access_peer", return_value=False), + ): + reason = FlashInferCustomAllReduce.unsupported_reason(group) + + assert reason is not None + assert "peer access" in reason + + +def test_custom_all_reduce_support_rejects_non_sm90_rank(): + group = ParallelGroup(object(), (0, 1), 0, 2) + with ( + patch.object(dist, "get_backend", return_value=dist.Backend.NCCL), + patch.object(torch.cuda, "is_available", return_value=True), + patch.object(torch.cuda, "device_count", return_value=2), + patch.object(torch.cuda, "current_device", return_value=0), + patch.object(torch.cuda, "get_device_capability", return_value=(8, 9)), + ): + reason = FlashInferCustomAllReduce.unsupported_reason(group) + + assert reason == "requires CUDA SM90, got (8, 9)" + + def test_parallel_context_lifecycle_and_local_groups(): reset_parallel_context() fake_groups = [] diff --git a/tests/test_sgl_moe.py b/tests/test_sgl_moe.py index 25624573..e8f10239 100644 --- a/tests/test_sgl_moe.py +++ b/tests/test_sgl_moe.py @@ -12,8 +12,10 @@ sgl_fused_moe_gate_support, sgl_moe_alignment_support, sgl_moe_align_block_size, + sgl_moe_ep_alignment_support, ) from sparsevllm.triton_kernel.moe import moe_align_block_size +from sparsevllm.triton_kernel.moe import fused_moe def test_sgl_moe_version_parser_accepts_post_release() -> None: @@ -40,11 +42,30 @@ def package_version(distribution: str) -> str: ): gate_supported, gate_reason = sgl_fused_moe_gate_support() align_supported, align_reason = sgl_moe_alignment_support() + ep_supported, ep_reason = sgl_moe_ep_alignment_support() assert gate_supported assert align_supported + assert ep_supported assert "0.4.5" in gate_reason assert "0.4.5" in align_reason + assert "0.4.5" in ep_reason + + +def test_sgl_moe_ep_support_accepts_canonical_03_package() -> None: + def package_version(distribution: str) -> str: + if distribution == "sgl-kernel": + return "0.3.21" + raise importlib.metadata.PackageNotFoundError(distribution) + + with ( + patch("importlib.util.find_spec", return_value=object()), + patch("importlib.metadata.version", side_effect=package_version), + ): + supported, reason = sgl_moe_ep_alignment_support() + + assert supported + assert "0.3.21" in reason @pytest.mark.skipif( @@ -152,3 +173,154 @@ def test_sgl_moe_alignment_matches_grouped_assignments() -> None: if value < num_assignments ) assert actual_ids == reference_ids + + +@pytest.mark.skipif( + not torch.cuda.is_available() or not sgl_moe_ep_alignment_support()[0], + reason="CUDA and a validated sgl-kernel are required", +) +@pytest.mark.parametrize(("local_start", "local_end"), [(0, 32), (32, 64)]) +def test_sgl_moe_alignment_matches_ep_shard( + local_start: int, + local_end: int, +) -> None: + torch.manual_seed(20260810) + topk_ids = torch.randint( + 0, + 64, + (32, 4), + dtype=torch.int32, + device="cuda", + ) + topk_ids[0] = torch.tensor([0, 31, 32, 63], device="cuda") + reference_sorted, reference_experts, reference_count = moe_align_block_size( + topk_ids, + 16, + 64, + local_expert_start=local_start, + local_expert_end=local_end, + ) + actual = sgl_moe_align_block_size( + topk_ids, + block_size=16, + num_experts=64, + local_expert_start=local_start, + local_expert_end=local_end, + ) + torch.cuda.synchronize() + num_assignments = int(topk_ids.numel()) + + def grouped_assignments(sorted_ids, expert_ids, count): + grouped = {} + for block_index in range(int(count.item()) // 16): + expert_id = int(expert_ids[block_index].item()) + if expert_id < 0: + continue + block = slice(block_index * 16, (block_index + 1) * 16) + grouped.setdefault(expert_id, []).extend( + value + for value in sorted_ids[block].tolist() + if value < num_assignments + ) + return { + expert_id: sorted(assignment_ids) + for expert_id, assignment_ids in grouped.items() + } + + assert grouped_assignments( + actual.sorted_token_ids, + actual.expert_ids, + actual.num_tokens_post_padded, + ) == grouped_assignments( + reference_sorted, + reference_experts, + reference_count, + ) + + +@pytest.mark.skipif( + not torch.cuda.is_available() or not sgl_moe_ep_alignment_support()[0], + reason="CUDA and a validated sgl-kernel are required", +) +@pytest.mark.parametrize("num_tokens", [5, 64]) +@pytest.mark.parametrize("all_remote", [False, True]) +def test_sgl_ep_alignment_preserves_full_fused_moe_output( + num_tokens: int, + all_remote: bool, +) -> None: + torch.manual_seed(20260810 + num_tokens + int(all_remote)) + device = torch.device("cuda") + hidden_size = 64 + intermediate_size = 32 + num_experts = 8 + num_local_experts = 4 + top_k = 2 + hidden_states = torch.randn( + num_tokens, + hidden_size, + dtype=torch.bfloat16, + device=device, + ) + w13_weight = torch.randn( + num_local_experts, + 2 * intermediate_size, + hidden_size, + dtype=torch.bfloat16, + device=device, + ) + w2_weight = torch.randn( + num_local_experts, + hidden_size, + intermediate_size, + dtype=torch.bfloat16, + device=device, + ) + if all_remote: + topk_ids = torch.randint( + num_local_experts, + num_experts, + (num_tokens, top_k), + dtype=torch.int64, + device=device, + ) + else: + topk_ids = torch.randint( + 0, + num_experts, + (num_tokens, top_k), + dtype=torch.int64, + device=device, + ) + topk_ids[0] = torch.tensor([0, num_local_experts], device=device) + topk_weights = torch.rand( + num_tokens, + top_k, + dtype=torch.bfloat16, + device=device, + ) + topk_weights /= topk_weights.sum(dim=-1, keepdim=True) + kwargs = { + "num_experts": num_experts, + "local_expert_start": 0, + } + + expected = fused_moe( + hidden_states, + w13_weight, + w2_weight, + topk_ids, + topk_weights, + **kwargs, + ) + actual = fused_moe( + hidden_states, + w13_weight, + w2_weight, + topk_ids, + topk_weights, + alignment_impl=sgl_moe_align_block_size, + **kwargs, + ) + torch.cuda.synchronize() + + assert torch.equal(actual, expected) From 4ba9354d6dc36ff769cea210526b9f3a45f10de5 Mon Sep 17 00:00:00 2001 From: CURRENTF <1206013760@qq.com> Date: Wed, 12 Aug 2026 15:03:24 +0800 Subject: [PATCH 22/48] perf: fuse glm moe routing --- src/sparsevllm/models/glm4_moe_lite.py | 34 ++-- src/sparsevllm/operators/moe.py | 5 +- src/sparsevllm/operators/moe_router.py | 52 ++++++ .../triton_kernel/moe_biased_sigmoid.py | 110 +++++++++++- tests/test_glm4_moe_lite.py | 14 +- tests/test_moe_router.py | 166 ++++++++++++++++++ tests/test_operator_providers.py | 2 +- tests/test_sgl_moe.py | 67 ++++++- 8 files changed, 415 insertions(+), 35 deletions(-) create mode 100644 tests/test_moe_router.py diff --git a/src/sparsevllm/models/glm4_moe_lite.py b/src/sparsevllm/models/glm4_moe_lite.py index 50b4b7a6..c71adca8 100644 --- a/src/sparsevllm/models/glm4_moe_lite.py +++ b/src/sparsevllm/models/glm4_moe_lite.py @@ -36,6 +36,10 @@ model_activation_dtype, resolve_moe_provider, ) +from sparsevllm.operators.moe_router import ( + MoeRouterOpSpec, + resolve_moe_router_provider, +) from sparsevllm.platforms import device_runtime from sparsevllm.utils.context import get_context from sparsevllm.utils.log import logger @@ -400,33 +404,27 @@ def __init__(self, config: Glm4MoeLiteConfig) -> None: self.e_score_correction_bias = nn.Parameter( torch.empty(self.num_experts, dtype=torch.float32) ) - from sparsevllm.operators.sgl_moe import SglGlmFusedMoeGate - - self.topk_impl = SglGlmFusedMoeGate( + self.op_spec = MoeRouterOpSpec( num_experts=self.num_experts, top_k=self.top_k, + activation_dtype=torch.float32, + norm_topk_prob=True, + cuda_graph=bool(getattr(config, "decode_cuda_graph", False)), + routing_method="biased_sigmoid", ) - self._fused_topk_impl = self.topk_impl + self.provider = resolve_moe_router_provider(self.op_spec) def forward( self, hidden_states: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: router_logits = F.linear(hidden_states.float(), self.weight) - if self.topk_impl is self._fused_topk_impl: - topk_weights, topk_ids = self.topk_impl( - router_logits, - self.e_score_correction_bias, - top_k=self.top_k, - routed_scaling_factor=self.routed_scaling_factor, - ) - else: - topk_weights, topk_ids = self.topk_impl( - router_logits, - self.e_score_correction_bias, - top_k=self.top_k, - ) - topk_weights = topk_weights * self.routed_scaling_factor + topk_weights, topk_ids = self.provider.run( + self.op_spec, + router_logits, + self.e_score_correction_bias, + routed_scaling_factor=self.routed_scaling_factor, + ) return router_logits, topk_weights, topk_ids diff --git a/src/sparsevllm/operators/moe.py b/src/sparsevllm/operators/moe.py index 3ee15e7b..8499e9a3 100644 --- a/src/sparsevllm/operators/moe.py +++ b/src/sparsevllm/operators/moe.py @@ -246,7 +246,10 @@ def run( num_tokens = int(hidden_states.shape[0]) if ( num_tokens <= 64 - and num_tokens * int(spec.top_k) * 4 > int(spec.num_experts) + and ( + int(spec.ep_size) > 1 + or num_tokens * int(spec.top_k) * 4 > int(spec.num_experts) + ) ): alignment_impl = sgl_moe_align_block_size return fused_moe( diff --git a/src/sparsevllm/operators/moe_router.py b/src/sparsevllm/operators/moe_router.py index fc8e74c3..148106ec 100644 --- a/src/sparsevllm/operators/moe_router.py +++ b/src/sparsevllm/operators/moe_router.py @@ -20,6 +20,7 @@ class MoeRouterOpSpec: activation_dtype: torch.dtype norm_topk_prob: bool cuda_graph: bool + routing_method: str = "softmax" def __post_init__(self) -> None: if self.num_experts <= 0: @@ -34,6 +35,8 @@ def __post_init__(self) -> None: "MoE router activations must be floating point, " f"got {self.activation_dtype}." ) + if self.routing_method not in {"softmax", "biased_sigmoid"}: + raise ValueError(f"Unsupported MoE routing method {self.routing_method!r}.") class MoeRouterProvider: @@ -44,6 +47,9 @@ def run( self, spec: MoeRouterOpSpec, router_logits: torch.Tensor, + correction_bias: torch.Tensor | None = None, + *, + routed_scaling_factor: float = 1.0, ) -> tuple[torch.Tensor, torch.Tensor]: raise NotImplementedError @@ -64,6 +70,8 @@ def supports( spec: MoeRouterOpSpec, caps: DeviceCaps, ) -> SupportResult: + if spec.routing_method != "softmax": + return SupportResult.no("requires softmax routing") if caps.platform != PlatformEnum.CUDA: return SupportResult.no(f"requires CUDA, got {caps.platform.name}") if not caps.supports_triton: @@ -84,7 +92,12 @@ def run( self, spec: MoeRouterOpSpec, router_logits: torch.Tensor, + correction_bias: torch.Tensor | None = None, + *, + routed_scaling_factor: float = 1.0, ) -> tuple[torch.Tensor, torch.Tensor]: + if correction_bias is not None or routed_scaling_factor != 1.0: + raise ValueError("Softmax routing does not accept bias or route scaling.") from sparsevllm.triton_kernel.moe_topk import topk_softmax return topk_softmax( @@ -94,6 +107,45 @@ def run( ) +@MOE_ROUTER_REGISTRY.register +class GlmBiasedSigmoidRouterProvider(MoeRouterProvider): + name = "triton_glm_biased_sigmoid" + priority = 20 + + @classmethod + def supports(cls, spec: MoeRouterOpSpec, caps: DeviceCaps) -> SupportResult: + if spec.routing_method != "biased_sigmoid": + return SupportResult.no("requires biased-sigmoid routing") + if caps.platform != PlatformEnum.CUDA or not caps.supports_triton: + return SupportResult.no("requires CUDA with Triton") + if spec.cuda_graph and not caps.supports_graph_capture: + return SupportResult.no("device does not support CUDA Graph capture") + if (spec.num_experts, spec.top_k) != (64, 4): + return SupportResult.no("requires 64 experts and top-k 4") + return SupportResult.yes() + + def run( + self, + spec: MoeRouterOpSpec, + router_logits: torch.Tensor, + correction_bias: torch.Tensor | None = None, + *, + routed_scaling_factor: float = 1.0, + ) -> tuple[torch.Tensor, torch.Tensor]: + if correction_bias is None: + raise ValueError("Biased-sigmoid routing requires correction_bias.") + from sparsevllm.triton_kernel.moe_biased_sigmoid import ( + fused_topk_biased_sigmoid, + ) + + return fused_topk_biased_sigmoid( + router_logits, + correction_bias, + top_k=spec.top_k, + routed_scaling_factor=routed_scaling_factor, + ) + + def resolve_moe_router_provider( spec: MoeRouterOpSpec, *, diff --git a/src/sparsevllm/triton_kernel/moe_biased_sigmoid.py b/src/sparsevllm/triton_kernel/moe_biased_sigmoid.py index 00ef78d5..d141de3b 100644 --- a/src/sparsevllm/triton_kernel/moe_biased_sigmoid.py +++ b/src/sparsevllm/triton_kernel/moe_biased_sigmoid.py @@ -8,6 +8,66 @@ _SUPPORTED_SHAPES = frozenset({(64, 4), (256, 8)}) +@triton.jit +def _fused_topk_biased_sigmoid_kernel( + logits_ptr, + correction_bias_ptr, + weights_ptr, + ids_ptr, + stride_logits_m, + stride_weights_m, + stride_ids_m, + normalization_epsilon: tl.constexpr, + routed_scaling_factor: tl.constexpr, + NUM_EXPERTS: tl.constexpr, + TOP_K: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + row = tl.program_id(0) + offsets = tl.arange(0, BLOCK_SIZE) + expert_mask = offsets < NUM_EXPERTS + logits = tl.load( + logits_ptr + row * stride_logits_m + offsets, + mask=expert_mask, + other=-float("inf"), + ) + routing_weights = tl.sigmoid(logits) + correction_bias = tl.load( + correction_bias_ptr + offsets, + mask=expert_mask, + other=0.0, + ) + scores = routing_weights + correction_bias + selection_values = tl.where(scores == scores, scores, float("inf")) + threshold = tl.min(tl.topk(selection_values, TOP_K), axis=0) + greater_mask = expert_mask & (selection_values > threshold) + equal_mask = expert_mask & (selection_values == threshold) + greater_rank = tl.cumsum(greater_mask.to(tl.int32), axis=0) - 1 + equal_rank = tl.cumsum(equal_mask.to(tl.int32), axis=0) - 1 + num_greater = tl.sum(greater_mask.to(tl.int32), axis=0) + selected_equal = equal_mask & (equal_rank < TOP_K - num_greater) + selected = greater_mask | selected_equal + output_slot = tl.where( + greater_mask, + greater_rank, + num_greater + equal_rank, + ) + denominator = tl.sum( + tl.where(selected, routing_weights, 0.0), + axis=0, + ) + weights_base = weights_ptr + row * stride_weights_m + ids_base = ids_ptr + row * stride_ids_m + tl.store( + weights_base + output_slot, + routing_weights + / (denominator + normalization_epsilon) + * routed_scaling_factor, + mask=selected, + ) + tl.store(ids_base + output_slot, offsets, mask=selected) + + @triton.jit def _topk_biased_sigmoid_kernel( routing_weights_ptr, @@ -91,6 +151,54 @@ def _validate_inputs( return num_tokens, num_experts +def fused_topk_biased_sigmoid( + router_logits: torch.Tensor, + correction_bias: torch.Tensor, + *, + top_k: int, + routed_scaling_factor: float, + normalization_epsilon: float = 1e-20, +) -> tuple[torch.Tensor, torch.Tensor]: + """Route the fixed GLM 64x4 shape in one graph-capturable kernel.""" + + num_tokens, num_experts = _validate_inputs( + router_logits, + correction_bias, + top_k=top_k, + ) + if (num_experts, int(top_k)) != (64, 4): + raise ValueError( + "Fused biased-sigmoid routing requires 64 experts and top-k 4, " + f"got {num_experts} and {top_k}." + ) + weights = torch.empty( + (num_tokens, int(top_k)), + dtype=torch.float32, + device=router_logits.device, + ) + ids = torch.empty( + (num_tokens, int(top_k)), + dtype=torch.int32, + device=router_logits.device, + ) + _fused_topk_biased_sigmoid_kernel[(num_tokens,)]( + router_logits, + correction_bias, + weights, + ids, + router_logits.stride(0), + weights.stride(0), + ids.stride(0), + normalization_epsilon=float(normalization_epsilon), + routed_scaling_factor=float(routed_scaling_factor), + NUM_EXPERTS=num_experts, + TOP_K=int(top_k), + BLOCK_SIZE=triton.next_power_of_2(num_experts), + num_warps=1, + ) + return weights, ids + + def topk_biased_sigmoid( router_logits: torch.Tensor, correction_bias: torch.Tensor, @@ -130,4 +238,4 @@ def topk_biased_sigmoid( return weights, ids -__all__ = ["topk_biased_sigmoid"] +__all__ = ["fused_topk_biased_sigmoid", "topk_biased_sigmoid"] diff --git a/tests/test_glm4_moe_lite.py b/tests/test_glm4_moe_lite.py index bc3e29c7..bac09c85 100644 --- a/tests/test_glm4_moe_lite.py +++ b/tests/test_glm4_moe_lite.py @@ -327,18 +327,14 @@ def test_glm_router_uses_bias_only_for_selection_and_scales_weights() -> None: router.e_score_correction_bias.data.zero_() router.e_score_correction_bias.data[:4] = 100.0 - def torch_topk(logits, correction_bias, *, top_k): + def torch_topk(_spec, logits, correction_bias, *, routed_scaling_factor): scores = logits.sigmoid() - ids = torch.topk( - scores + correction_bias, - top_k, - dim=-1, - sorted=False, - ).indices + ids = torch.topk(scores + correction_bias, 4, dim=-1, sorted=False).indices weights = scores.gather(1, ids) - return weights / (weights.sum(dim=-1, keepdim=True) + 1e-20), ids + weights = weights / (weights.sum(dim=-1, keepdim=True) + 1e-20) + return weights * routed_scaling_factor, ids - router.topk_impl = torch_topk + router.provider.run = torch_topk hidden_states = torch.randn(7, 64, dtype=torch.bfloat16) logits, weights, ids = router(hidden_states) original_scores = logits.sigmoid() diff --git a/tests/test_moe_router.py b/tests/test_moe_router.py new file mode 100644 index 00000000..c37432f3 --- /dev/null +++ b/tests/test_moe_router.py @@ -0,0 +1,166 @@ +from __future__ import annotations + +import pytest +import torch + +from sparsevllm.operators.moe_router import ( + GlmBiasedSigmoidRouterProvider, + MoeRouterOpSpec, +) + + +def _router(): + spec = MoeRouterOpSpec(64, 4, torch.float32, True, True, "biased_sigmoid") + return spec, GlmBiasedSigmoidRouterProvider() + + +def _reference_routes( + logits: torch.Tensor, + correction_bias: torch.Tensor, + scaling: float, +) -> tuple[torch.Tensor, torch.Tensor]: + routing_weights = torch.sigmoid(logits) + ids = torch.topk( + routing_weights + correction_bias, + 4, + dim=-1, + sorted=False, + ).indices + weights = routing_weights.gather(1, ids) + weights /= weights.sum(dim=-1, keepdim=True) + 1e-20 + return weights * scaling, ids + + +def _assert_same_routes( + actual_weights: torch.Tensor, + actual_ids: torch.Tensor, + reference_weights: torch.Tensor, + reference_ids: torch.Tensor, +) -> None: + reference_order = reference_ids.argsort(dim=-1) + actual_order = actual_ids.argsort(dim=-1) + assert torch.equal( + actual_ids.gather(1, actual_order), + reference_ids.gather(1, reference_order), + ) + torch.testing.assert_close( + actual_weights.gather(1, actual_order), + reference_weights.gather(1, reference_order), + rtol=2e-6, + atol=2e-7, + equal_nan=True, + ) + + +def test_glm_router_rejects_wrong_shape() -> None: + spec = MoeRouterOpSpec(32, 4, torch.float32, True, True, "biased_sigmoid") + from sparsevllm.platforms import DeviceCaps, PlatformEnum + + caps = DeviceCaps( + platform=PlatformEnum.CUDA, + device_type="cuda", + device_index=0, + device_name="test", + supports_triton=True, + ) + assert not GlmBiasedSigmoidRouterProvider.supports(spec, caps).supported + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize("num_tokens", [1, 2, 4, 32]) +def test_glm_router_matches_reference_and_replays_updated_graph( + num_tokens: int, +) -> None: + torch.manual_seed(20260810 + num_tokens) + logits = ( + torch.randn(num_tokens, 64, dtype=torch.float32, device="cuda") * 3 + ) + correction_bias = ( + torch.randn(64, dtype=torch.float32, device="cuda") * 0.1 + ) + scaling = 1.8 + reference_weights, reference_ids = _reference_routes( + logits, + correction_bias, + scaling, + ) + + spec, router = _router() + actual_weights, actual_ids = router.run( + spec, + logits, + correction_bias, + routed_scaling_factor=scaling, + ) + _assert_same_routes( + actual_weights, + actual_ids, + reference_weights, + reference_ids, + ) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_weights, graph_ids = router.run( + spec, + logits, + correction_bias, + routed_scaling_factor=scaling, + ) + + logits.copy_(torch.randn_like(logits) * 4) + correction_bias.copy_(torch.randn_like(correction_bias) * 0.2) + replay_reference_weights, replay_reference_ids = _reference_routes( + logits, + correction_bias, + scaling, + ) + graph.replay() + torch.cuda.synchronize() + _assert_same_routes( + graph_weights, + graph_ids, + replay_reference_weights, + replay_reference_ids, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_glm_router_handles_ties_extremes_and_nonfinite_scores() -> None: + logits = torch.linspace( + -80, + 80, + 64, + dtype=torch.float32, + device="cuda", + ).repeat(3, 1) + correction_bias = torch.zeros(64, dtype=torch.float32, device="cuda") + logits[0].zero_() + logits[2, 7] = float("nan") + scaling = 1.8 + spec, router = _router() + weights, ids = router.run( + spec, + logits, + correction_bias, + routed_scaling_factor=scaling, + ) + + assert torch.equal(ids[0].sort().values, torch.arange(4, device="cuda")) + torch.testing.assert_close( + weights[0], + torch.full((4,), scaling / 4, device="cuda"), + rtol=0, + atol=0, + ) + reference_weights, reference_ids = _reference_routes( + logits[1:], + correction_bias, + scaling, + ) + _assert_same_routes( + weights[1:], + ids[1:], + reference_weights, + reference_ids, + ) diff --git a/tests/test_operator_providers.py b/tests/test_operator_providers.py index 7b4a4537..98e52c7f 100644 --- a/tests/test_operator_providers.py +++ b/tests/test_operator_providers.py @@ -571,7 +571,7 @@ def test_glm_tp2_ep2_moe_uses_sgl_aligned_provider(): @pytest.mark.parametrize( ("num_tokens", "expects_sgl_alignment"), - [(4, False), (5, True), (64, True), (65, False)], + [(1, True), (2, True), (4, True), (5, True), (64, True), (65, False)], ) def test_glm_tp2_ep2_sgl_alignment_is_bounded( num_tokens, diff --git a/tests/test_sgl_moe.py b/tests/test_sgl_moe.py index e8f10239..d318fad4 100644 --- a/tests/test_sgl_moe.py +++ b/tests/test_sgl_moe.py @@ -242,11 +242,13 @@ def grouped_assignments(sorted_ids, expert_ids, count): not torch.cuda.is_available() or not sgl_moe_ep_alignment_support()[0], reason="CUDA and a validated sgl-kernel are required", ) -@pytest.mark.parametrize("num_tokens", [5, 64]) +@pytest.mark.parametrize("num_tokens", [1, 2, 4, 5, 64]) @pytest.mark.parametrize("all_remote", [False, True]) +@pytest.mark.parametrize("local_start", [0, 4]) def test_sgl_ep_alignment_preserves_full_fused_moe_output( num_tokens: int, all_remote: bool, + local_start: int, ) -> None: torch.manual_seed(20260810 + num_tokens + int(all_remote)) device = torch.device("cuda") @@ -276,9 +278,10 @@ def test_sgl_ep_alignment_preserves_full_fused_moe_output( device=device, ) if all_remote: + remote_start = num_local_experts if local_start == 0 else 0 topk_ids = torch.randint( - num_local_experts, - num_experts, + remote_start, + remote_start + num_local_experts, (num_tokens, top_k), dtype=torch.int64, device=device, @@ -291,7 +294,11 @@ def test_sgl_ep_alignment_preserves_full_fused_moe_output( dtype=torch.int64, device=device, ) - topk_ids[0] = torch.tensor([0, num_local_experts], device=device) + remote_id = num_local_experts if local_start == 0 else 0 + topk_ids[0] = torch.tensor( + [local_start, remote_id], + device=device, + ) topk_weights = torch.rand( num_tokens, top_k, @@ -301,7 +308,7 @@ def test_sgl_ep_alignment_preserves_full_fused_moe_output( topk_weights /= topk_weights.sum(dim=-1, keepdim=True) kwargs = { "num_experts": num_experts, - "local_expert_start": 0, + "local_expert_start": local_start, } expected = fused_moe( @@ -324,3 +331,53 @@ def test_sgl_ep_alignment_preserves_full_fused_moe_output( torch.cuda.synchronize() assert torch.equal(actual, expected) + + if num_tokens == 2: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_actual = fused_moe( + hidden_states, + w13_weight, + w2_weight, + topk_ids, + topk_weights, + alignment_impl=sgl_moe_align_block_size, + **kwargs, + ) + hidden_states.copy_(torch.randn_like(hidden_states)) + if all_remote: + replay_ids = torch.randint( + remote_start, + remote_start + num_local_experts, + topk_ids.shape, + dtype=topk_ids.dtype, + device=device, + ) + else: + replay_ids = torch.randint( + 0, + num_experts, + topk_ids.shape, + dtype=topk_ids.dtype, + device=device, + ) + remote_id = num_local_experts if local_start == 0 else 0 + replay_ids[0] = torch.tensor( + [local_start + 1, remote_id], + device=device, + ) + topk_ids.copy_(replay_ids) + replay_weights = torch.rand_like(topk_weights) + replay_weights /= replay_weights.sum(dim=-1, keepdim=True) + topk_weights.copy_(replay_weights) + replay_expected = fused_moe( + hidden_states, + w13_weight, + w2_weight, + topk_ids, + topk_weights, + **kwargs, + ) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(graph_actual, replay_expected) From 5639c650a4a9c61e84797d977463282327190697 Mon Sep 17 00:00:00 2001 From: CURRENTF <1206013760@qq.com> Date: Sun, 9 Aug 2026 17:57:08 +0800 Subject: [PATCH 23/48] perf: accelerate minimax long context --- .../distributed/parallel_context.py | 29 +- src/sparsevllm/engine/model_runner.py | 6 + src/sparsevllm/layers/attention.py | 54 ++- src/sparsevllm/layers/attention_backend.py | 54 ++- src/sparsevllm/models/minimax_m2.py | 26 +- src/sparsevllm/operators/all_reduce.py | 230 +++++++---- src/sparsevllm/operators/moe.py | 8 +- src/sparsevllm/operators/prefill_attention.py | 357 ++++++++++++++++++ .../gqa_flash_decoding_stage1.py | 20 +- tests/test_operator_providers.py | 37 +- tests/test_parallel_context.py | 18 - tests/test_prefill_attention_provider.py | 164 ++++++++ tests/test_qwen35_hd256_decode_routing.py | 37 ++ 13 files changed, 906 insertions(+), 134 deletions(-) create mode 100644 src/sparsevllm/operators/prefill_attention.py create mode 100644 tests/test_prefill_attention_provider.py diff --git a/src/sparsevllm/distributed/parallel_context.py b/src/sparsevllm/distributed/parallel_context.py index 43a4c6b2..262c4227 100644 --- a/src/sparsevllm/distributed/parallel_context.py +++ b/src/sparsevllm/distributed/parallel_context.py @@ -1,14 +1,13 @@ from __future__ import annotations from contextlib import contextmanager -from dataclasses import dataclass, field, replace +from dataclasses import dataclass, replace from typing import Iterator, Protocol import torch import torch.distributed as dist from sparsevllm.distributed.topology import ParallelTopology, parallel_group_ranks -from sparsevllm.operators.all_reduce import AllReduceProvider, resolve_all_reduce_provider from sparsevllm.utils.log import logger @@ -28,7 +27,6 @@ class ParallelGroup: ranks: tuple[int, ...] rank: int size: int - all_reduce_provider: AllReduceProvider | None = field(default=None, compare=False, repr=False) def __post_init__(self) -> None: if self.size != len(self.ranks): @@ -109,10 +107,7 @@ def _all_reduce( op: dist.ReduceOp = dist.ReduceOp.SUM, ) -> torch.Tensor: if group.size > 1: - if op != dist.ReduceOp.SUM or group.all_reduce_provider is None: - dist.all_reduce(tensor, op=op, group=group.process_group) - else: - tensor = group.all_reduce_provider.run(tensor) + dist.all_reduce(tensor, op=op, group=group.process_group) return tensor def world_all_reduce( @@ -307,25 +302,7 @@ def init_parallel_context( ranks_by_dimension["moe_tensor"], process_groups, world_rank ), ) - providers: dict[tuple[int, ...], AllReduceProvider] = {} - - def bind_provider(group: ParallelGroup) -> ParallelGroup: - if group.size == 1: - return group - provider = providers.get(group.ranks) - if provider is None: - provider = providers[group.ranks] = resolve_all_reduce_provider( - group.process_group, group.size - ) - return replace(group, all_reduce_provider=provider) - - _PARALLEL_CONTEXT = ParallelContext( - world=bind_provider(context.world), - tensor=bind_provider(context.tensor), - expert=bind_provider(context.expert), - data=bind_provider(context.data), - moe_tensor=bind_provider(context.moe_tensor or context.tensor), - ) + _PARALLEL_CONTEXT = context if enable_flashinfer_custom_all_reduce: from sparsevllm.distributed.flashinfer_custom_all_reduce import ( FlashInferCustomAllReduce, diff --git a/src/sparsevllm/engine/model_runner.py b/src/sparsevllm/engine/model_runner.py index 565340e8..bd6bd211 100644 --- a/src/sparsevllm/engine/model_runner.py +++ b/src/sparsevllm/engine/model_runner.py @@ -170,6 +170,12 @@ def __init__( torch.set_default_dtype(hf_config.torch_dtype) torch.set_default_device(self.device) setattr(hf_config, "mlp_chunk_size", config.mlp_chunk_size) + if hf_config.model_type == "minimax_m2": + setattr( + hf_config, + "prefill_kv_view_layer_invariant", + config.vllm_sparse_method == "", + ) setattr( hf_config, "decode_cuda_graph", diff --git a/src/sparsevllm/layers/attention.py b/src/sparsevllm/layers/attention.py index 4af57348..6f84e48a 100644 --- a/src/sparsevllm/layers/attention.py +++ b/src/sparsevllm/layers/attention.py @@ -5,6 +5,10 @@ from sparsevllm.engine.cache_manager import ExplicitKVPayload from sparsevllm.layers.attention_backend import TritonAttentionBackend +from sparsevllm.operators.prefill_attention import ( + PrefillAttentionOpSpec, + resolve_prefill_attention_provider, +) from sparsevllm.utils.context import get_context from sparsevllm.engine.sparse_controller import SparseController @@ -57,6 +61,8 @@ def __init__( head_dim, scale, num_kv_heads, + *, + prefill_op_spec: PrefillAttentionOpSpec | None = None, ): super().__init__() self.num_heads = num_heads @@ -64,6 +70,14 @@ def __init__( self.scale = scale self.num_kv_heads = num_kv_heads self.attention_backend = TritonAttentionBackend() + self.prefill_op_spec = prefill_op_spec + self.prefill_provider = ( + resolve_prefill_attention_provider(prefill_op_spec) + if prefill_op_spec is not None + else None + ) + if self.prefill_provider is not None: + self.prefill_provider.prepare(prefill_op_spec) def forward( self, @@ -108,13 +122,37 @@ def forward( else: max_input_len = prefill_meta.context_lens.max().item() - o = self.attention_backend.run_prefill( + fake_output = self.attention_backend.maybe_run_fake_prefill( q, prefill_view, - b_start_loc=b_start_loc, chunk_lens=chunk_lens, max_input_len=max_input_len, ) + if fake_output is not None: + o = fake_output + elif self.prefill_provider is None: + o = self.attention_backend.run_prefill( + q, + prefill_view, + b_start_loc=b_start_loc, + chunk_lens=chunk_lens, + max_input_len=max_input_len, + ) + else: + self.attention_backend.debug_check_prefill_bounds( + q, + prefill_view, + chunk_lens=chunk_lens, + ) + o = self.prefill_provider.run( + self.prefill_op_spec, + q, + prefill_view, + qo_indptr=context.cu_seqlens_q, + chunk_lens=chunk_lens, + max_context_len=max_input_len, + layer_idx=int(layer_idx), + ) cache_manager.collect_prefill_attention_score( layer_idx, q, @@ -184,6 +222,16 @@ def forward( f"decode requires a positive context length, got {max_len_in_batch} at layer={layer_idx}" ) BLOCK_SEQ = cache_manager.get_decode_block_seq(layer_idx, 256) + BLOCK_SEQ, gqa_block_n, gqa_num_warps = ( + self.attention_backend.gqa_decode_launch_config( + block_seq=BLOCK_SEQ, + max_context_len=max_len_in_batch, + num_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + head_dim=self.head_dim, + requires_attention_scores=decode_meta.attn_score is not None, + ) + ) num_seq_blocks = (max_len_in_batch + BLOCK_SEQ - 1) // BLOCK_SEQ mid_o, mid_o_logexpsum = get_decode_workspace( @@ -204,6 +252,8 @@ def forward( block_seq=BLOCK_SEQ, num_heads=self.num_heads, num_kv_heads=self.num_kv_heads, + gqa_block_n=gqa_block_n, + gqa_num_warps=gqa_num_warps, ) cache_manager.record_decode_query(layer_idx, q) diff --git a/src/sparsevllm/layers/attention_backend.py b/src/sparsevllm/layers/attention_backend.py index 959ee27a..cb3ecd3e 100644 --- a/src/sparsevllm/layers/attention_backend.py +++ b/src/sparsevllm/layers/attention_backend.py @@ -118,6 +118,52 @@ class TritonAttentionBackend: def __init__(self) -> None: record_operator_binding("Attention", self) + @staticmethod + def gqa_decode_launch_config( + *, + block_seq: int, + max_context_len: int, + num_heads: int, + num_kv_heads: int, + head_dim: int, + requires_attention_scores: bool, + ) -> tuple[int, int, int]: + if ( + block_seq == 256 + and max_context_len > 32768 + and (num_heads, num_kv_heads, head_dim) == (12, 2, 128) + and not requires_attention_scores + ): + return 1024, 128, 4 + return block_seq, 16, 2 + + def maybe_run_fake_prefill( + self, + q: torch.Tensor, + view: PrefillComputeView, + *, + chunk_lens: torch.Tensor, + max_input_len: int, + ) -> torch.Tensor | None: + if not _fake_prefill_attention_enabled(): + return None + meta = view.meta + probe_context_tokens = int(meta.max_context_len or max_input_len) + real_probe_min_context = _warmup_real_prefill_probe_min_context() + if ( + real_probe_min_context is not None + and probe_context_tokens >= real_probe_min_context + ): + log_once( + "Warmup real prefill attention probe executing at " + f"context_tokens={probe_context_tokens} query_tokens={int(q.shape[0])} " + f"batch_seqs={int(chunk_lens.shape[0])}.", + level="INFO", + ) + return None + _fill_fake_attention_score(meta.attn_score) + return _fake_attention_output(q) + def run_prefill( self, q: torch.Tensor, @@ -140,7 +186,7 @@ def run_prefill( f"active_slots_shape={tuple(meta.active_slots.shape)}" ) b_prompt_cache_len = b_seq_len - chunk_lens - self._debug_check_prefill_bounds(q, view, chunk_lens=chunk_lens) + self.debug_check_prefill_bounds(q, view, chunk_lens=chunk_lens) if _fake_prefill_attention_enabled(): real_probe_min_context = _warmup_real_prefill_probe_min_context() probe_context_tokens = ( @@ -177,7 +223,7 @@ def run_prefill( ) return o - def _debug_check_prefill_bounds( + def debug_check_prefill_bounds( self, q: torch.Tensor, view: PrefillComputeView, @@ -250,6 +296,8 @@ def run_decode( block_seq: int, num_heads: int, num_kv_heads: int, + gqa_block_n: int = 16, + gqa_num_warps: int = 2, ) -> torch.Tensor: payload = _require_explicit_payload(view, operation="Triton decode") meta = view.meta @@ -349,6 +397,8 @@ def run_decode( mid_o, mid_o_logexpsum, block_seq, + gqa_block_n, + gqa_num_warps, ) else: mha_flash_decode_stage1( diff --git a/src/sparsevllm/models/minimax_m2.py b/src/sparsevllm/models/minimax_m2.py index 65cd7404..b07a8a1c 100644 --- a/src/sparsevllm/models/minimax_m2.py +++ b/src/sparsevllm/models/minimax_m2.py @@ -22,6 +22,7 @@ model_activation_dtype, resolve_moe_provider, ) +from sparsevllm.operators.prefill_attention import PrefillAttentionOpSpec from sparsevllm.platforms import device_runtime from sparsevllm.quantization.fp8_tp import Fp8ExpertTpShard from sparsevllm.utils.context import get_context @@ -336,11 +337,27 @@ def __init__(self, config) -> None: base=float(config.rope_theta), rope_scaling=None, ) + prefill_op_spec = None + if hasattr(config, "prefill_kv_view_layer_invariant"): + prefill_op_spec = PrefillAttentionOpSpec( + num_query_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + head_dim=self.head_dim, + activation_dtype=model_activation_dtype(config), + softmax_scale=self.head_dim**-0.5, + causal=True, + page_size=1, + requires_attention_scores=False, + layer_invariant_page_table=bool( + config.prefill_kv_view_layer_invariant + ), + ) self.attn = Attention( self.num_heads, self.head_dim, self.head_dim**-0.5, self.num_kv_heads, + prefill_op_spec=prefill_op_spec, ) def forward( @@ -641,11 +658,18 @@ def validate_loaded_weights(self, loaded_parameter_names: set[str]) -> None: f"weights={unexpected_skips[:4]}, scales={unexpected_scale_skips[:4]}." ) logger.info( - "Loaded MiniMax M2 rank {} provider={} attention TP {}/{} MoE TP " + "Loaded MiniMax M2 rank {} provider={} prefill_provider={} all_reduce_provider={} " + "attention TP {}/{} MoE TP " "{}/{} local experts [{}, {}) across {} layers; intentionally skipped " "{} remote expert weight/scale pairs.", self.parallel_context.world_rank, self.model.layers[0].block_sparse_moe.experts.provider.name, + getattr( + self.model.layers[0].self_attn.attn.prefill_provider, + "name", + "legacy_triton", + ), + getattr(self.parallel_context.all_reduce_provider, "name", "unconfigured"), self.parallel_context.tp_rank, self.parallel_context.tp_size, self.parallel_context.moe_tp_rank, diff --git a/src/sparsevllm/operators/all_reduce.py b/src/sparsevllm/operators/all_reduce.py index ee848e3d..153fce91 100644 --- a/src/sparsevllm/operators/all_reduce.py +++ b/src/sparsevllm/operators/all_reduce.py @@ -1,84 +1,180 @@ from __future__ import annotations -from typing import Protocol +import re +from dataclasses import dataclass +from importlib.metadata import PackageNotFoundError, version +from importlib.util import find_spec import torch import torch.distributed as dist -from sparsevllm.utils.log import logger - - -class AllReduceProvider(Protocol): - name: str - - def run(self, tensor: torch.Tensor) -> torch.Tensor: ... - - -class TorchDistributedAllReduceProvider: - name = "torch_distributed" - - def __init__(self, group: dist.ProcessGroup | None) -> None: - self.group = group - - def run(self, tensor: torch.Tensor) -> torch.Tensor: - dist.all_reduce(tensor, group=self.group) - return tensor - - -class HopperTp2FlashInferAllReduceProvider: - name = "hopper_tp2_flashinfer" - hidden_size = 2048 - max_rows = 256 - - def __init__(self, group: dist.ProcessGroup) -> None: - from flashinfer import comm - - self.comm = comm - self.fallback = TorchDistributedAllReduceProvider(group) - self.workspace = comm.create_allreduce_fusion_workspace( +import sparsevllm.platforms as platforms +from sparsevllm.operators.registry import ( + OpRegistry, + OpResolver, + SupportResult, + runtime_version_at_least, +) +from sparsevllm.platforms.interface import DeviceCaps, PlatformEnum + + +@dataclass(frozen=True) +class AllReduceOpSpec: + world_size: int + max_tokens: int + hidden_size: int + dtype: torch.dtype + cuda_graph: bool + + def __post_init__(self) -> None: + if self.world_size <= 0 or self.max_tokens <= 0 or self.hidden_size <= 0: + raise ValueError("All-reduce dimensions must be positive.") + + +class AllReduceProvider: + name = "" + priority = 0 + + def prepare( + self, + spec: AllReduceOpSpec, + *, + group: dist.ProcessGroup, + rank: int, + ) -> None: + del spec, group, rank + + def can_run(self, spec: AllReduceOpSpec, tensor: torch.Tensor) -> bool: + del spec, tensor + return True + + def run( + self, + spec: AllReduceOpSpec, + tensor: torch.Tensor, + *, + group: dist.ProcessGroup, + ) -> torch.Tensor: + raise NotImplementedError + + +ALL_REDUCE_REGISTRY: OpRegistry[AllReduceOpSpec, AllReduceProvider] = OpRegistry( + "all-reduce" +) + + +@ALL_REDUCE_REGISTRY.register +class FlashInferTrtllmAllReduceProvider(AllReduceProvider): + name = "flashinfer_trtllm_sm90" + priority = 100 + + @classmethod + def supports(cls, spec: AllReduceOpSpec, caps: DeviceCaps) -> SupportResult: + if caps.platform != PlatformEnum.CUDA or caps.compute_capability != (9, 0): + return SupportResult.no( + f"requires CUDA SM90, got {caps.platform.name} {caps.compute_capability}" + ) + if caps.device_name != "NVIDIA H100 80GB HBM3": + return SupportResult.no( + "requires profiled NVIDIA H100 80GB HBM3 hardware, " + f"got {caps.device_name}" + ) + if not runtime_version_at_least(caps.runtime_version, (12, 8)): + return SupportResult.no( + "requires CUDA runtime >= 12.8, " + f"got {caps.runtime_version or 'unknown'}" + ) + if not caps.supports_graph_capture or not spec.cuda_graph: + return SupportResult.no("requires CUDA Graph execution") + if spec.world_size != 4: + return SupportResult.no( + f"requires profiled world_size=4, got {spec.world_size}" + ) + if spec.max_tokens > 32: + return SupportResult.no( + f"requires max_tokens <= 32, got {spec.max_tokens}" + ) + if spec.hidden_size != 3072 or spec.dtype != torch.bfloat16: + return SupportResult.no( + "requires profiled BF16 hidden_size=3072, got " + f"{spec.dtype} hidden_size={spec.hidden_size}" + ) + if find_spec("flashinfer") is None: + return SupportResult.no("flashinfer is not installed") + try: + installed = version("flashinfer-python") + except PackageNotFoundError: + return SupportResult.no("flashinfer-python package metadata is unavailable") + numeric = tuple(int(part) for part in re.findall(r"\d+", installed)[:3]) + if numeric < (0, 6, 15): + return SupportResult.no( + f"requires flashinfer-python >= 0.6.15, got {installed}" + ) + return SupportResult.yes() + + def __init__(self) -> None: + self.workspace = None + + def prepare(self, spec, *, group, rank) -> None: + from flashinfer.comm import create_allreduce_fusion_workspace + + self.workspace = create_allreduce_fusion_workspace( backend="trtllm", - world_size=2, - rank=dist.get_rank(group), - max_token_num=self.max_rows, - hidden_dim=self.hidden_size, - dtype=torch.bfloat16, + world_size=spec.world_size, + rank=rank, + max_token_num=spec.max_tokens, + hidden_dim=spec.hidden_size, + dtype=spec.dtype, group=group, ) - def _supports(self, tensor: torch.Tensor) -> bool: + def can_run(self, spec, tensor) -> bool: return ( tensor.is_cuda - and tensor.dtype == torch.bfloat16 + and tensor.dtype == spec.dtype + and tensor.ndim == 2 + and 0 < int(tensor.shape[0]) <= spec.max_tokens + and int(tensor.shape[1]) == spec.hidden_size and tensor.is_contiguous() - and tensor.ndim >= 2 - and tensor.shape[-1] == self.hidden_size - and tensor.numel() <= self.max_rows * self.hidden_size ) - def run(self, tensor: torch.Tensor) -> torch.Tensor: - if not self._supports(tensor): - return self.fallback.run(tensor) - output = self.comm.allreduce_fusion( - input=tensor.view(-1, self.hidden_size), + def run(self, spec, tensor, *, group) -> torch.Tensor: + del group + if self.workspace is None: + raise RuntimeError("FlashInfer all-reduce provider was not prepared.") + if not self.can_run(spec, tensor): + raise ValueError( + "FlashInfer all-reduce received an unsupported tensor: " + f"shape={tuple(tensor.shape)} dtype={tensor.dtype} device={tensor.device}." + ) + from flashinfer.comm import allreduce_fusion + from flashinfer.comm.trtllm_ar import AllReduceFusionPattern + + return allreduce_fusion( + input=tensor, workspace=self.workspace, - pattern=self.comm.AllReduceFusionPattern.kAllReduce, - launch_with_pdl=True, - trigger_completion_at_end=tensor.numel() > 16 * self.hidden_size, + pattern=AllReduceFusionPattern.kAllReduce, ) - return output.view_as(tensor) - - -def resolve_all_reduce_provider( - group: dist.ProcessGroup | None, - world_size: int, -) -> AllReduceProvider: - if ( - world_size == 2 - and group is not None - and dist.get_backend(group) == dist.Backend.NCCL - and torch.cuda.get_device_capability() == (9, 0) - ): - provider = HopperTp2FlashInferAllReduceProvider(group) - logger.info("AllReduce provider: %s", provider.name) - return provider - return TorchDistributedAllReduceProvider(group) + + +@ALL_REDUCE_REGISTRY.register +class TorchDistributedAllReduceProvider(AllReduceProvider): + name = "torch_distributed" + priority = 10 + + @classmethod + def supports(cls, spec: AllReduceOpSpec, caps: DeviceCaps) -> SupportResult: + del spec, caps + return SupportResult.yes() + + def run(self, spec, tensor, *, group) -> torch.Tensor: + del spec + dist.all_reduce(tensor, group=group) + return tensor + + +def resolve_all_reduce_provider(spec: AllReduceOpSpec) -> AllReduceProvider: + platform = platforms.current_platform + device_index = torch.cuda.current_device() if platform.is_cuda_alike() else 0 + caps = platform.get_device_caps(int(device_index)) + return OpResolver(ALL_REDUCE_REGISTRY).resolve(spec, caps).provider diff --git a/src/sparsevllm/operators/moe.py b/src/sparsevllm/operators/moe.py index 8499e9a3..173e3598 100644 --- a/src/sparsevllm/operators/moe.py +++ b/src/sparsevllm/operators/moe.py @@ -267,7 +267,7 @@ def run( @MOE_REGISTRY.register class TritonMinimaxM2FusedMoeProvider(MoeProvider): name = "triton_minimax_m2_fused" - priority = 110 + priority = 125 gate_up_order = "gate_up" @classmethod @@ -369,7 +369,7 @@ def run( @MOE_REGISTRY.register class FlashInferCutlassFp8MoeProvider(MoeProvider): name = "flashinfer_cutlass_fp8_sm90" - priority = 100 + priority = 120 gate_up_order = "up_gate" @classmethod @@ -615,7 +615,7 @@ class HopperQwen36HybridFp8MoeProvider(FlashInferCutlassFp8MoeProvider): """Bind one weight layout and dispatch profiled token buckets by kernel.""" name = "hopper_qwen36_hybrid_fp8" - priority = 110 + priority = 130 PROFILED_DEVICE_NAME = "NVIDIA H100 80GB HBM3" PROFILED_SHAPES = frozenset( { @@ -711,7 +711,7 @@ def run( @MOE_REGISTRY.register class H20Qwen36HybridFp8MoeProvider(HopperQwen36HybridFp8MoeProvider): name = "h20_qwen36_hybrid_fp8" - priority = 111 + priority = 131 PROFILED_DEVICE_NAME = "NVIDIA H20" TRITON_MAX_TOKENS_BY_EP_SIZE = {1: 8, 2: 1} diff --git a/src/sparsevllm/operators/prefill_attention.py b/src/sparsevllm/operators/prefill_attention.py new file mode 100644 index 00000000..f2672b06 --- /dev/null +++ b/src/sparsevllm/operators/prefill_attention.py @@ -0,0 +1,357 @@ +from __future__ import annotations + +import re +from dataclasses import dataclass +from importlib.metadata import PackageNotFoundError, version +from importlib.util import find_spec +from typing import Any + +import torch + +import sparsevllm.platforms as platforms +from sparsevllm.operators.registry import ( + OpRegistry, + OpResolver, + SupportResult, + runtime_version_at_least, +) +from sparsevllm.platforms.interface import DeviceCaps, PlatformEnum + + +@dataclass(frozen=True) +class PrefillAttentionOpSpec: + num_query_heads: int + num_kv_heads: int + head_dim: int + activation_dtype: torch.dtype + softmax_scale: float + causal: bool = True + page_size: int = 1 + requires_attention_scores: bool = False + layer_invariant_page_table: bool = False + + def __post_init__(self) -> None: + if self.num_query_heads <= 0 or self.num_kv_heads <= 0: + raise ValueError("Prefill attention head counts must be positive.") + if self.num_query_heads % self.num_kv_heads: + raise ValueError("Query heads must be divisible by KV heads.") + if self.head_dim <= 0 or self.page_size <= 0: + raise ValueError("Prefill attention dimensions must be positive.") + if self.softmax_scale <= 0: + raise ValueError("Prefill attention softmax_scale must be positive.") + + +class PrefillAttentionProvider: + name = "" + priority = 0 + + def prepare(self, spec: PrefillAttentionOpSpec) -> None: + del spec + + def run( + self, + spec: PrefillAttentionOpSpec, + q: torch.Tensor, + view: Any, + *, + qo_indptr: torch.Tensor, + chunk_lens: torch.Tensor, + max_context_len: int, + layer_idx: int, + ) -> torch.Tensor: + raise NotImplementedError + + +PREFILL_ATTENTION_REGISTRY: OpRegistry[ + PrefillAttentionOpSpec, PrefillAttentionProvider +] = OpRegistry("paged prefill attention") + + +def _view_parts(view: Any) -> tuple[Any, Any]: + """Return the physical payload and logical metadata for a prefill view.""" + + return getattr(view, "payload", view), getattr(view, "meta", view) + + +@PREFILL_ATTENTION_REGISTRY.register +class FlashInferPagedPrefillAttentionProvider(PrefillAttentionProvider): + name = "flashinfer_paged_prefill_fa3_sm90" + priority = 100 + + @classmethod + def supports( + cls, spec: PrefillAttentionOpSpec, caps: DeviceCaps + ) -> SupportResult: + if caps.platform != PlatformEnum.CUDA or caps.compute_capability != (9, 0): + return SupportResult.no( + f"requires CUDA SM90, got {caps.platform.name} {caps.compute_capability}" + ) + if caps.device_name != "NVIDIA H100 80GB HBM3": + return SupportResult.no( + "requires profiled NVIDIA H100 80GB HBM3 hardware, " + f"got {caps.device_name}" + ) + if not runtime_version_at_least(caps.runtime_version, (12, 8)): + return SupportResult.no( + "requires CUDA runtime >= 12.8, " + f"got {caps.runtime_version or 'unknown'}" + ) + if spec.activation_dtype != torch.bfloat16: + return SupportResult.no( + f"requires BF16 Q/K/V, got {spec.activation_dtype}" + ) + expected_shape = (12, 2, 128) + actual_shape = ( + spec.num_query_heads, + spec.num_kv_heads, + spec.head_dim, + ) + if actual_shape != expected_shape: + return SupportResult.no( + f"requires profiled local Q/KV/head shape {expected_shape}, got {actual_shape}" + ) + if not spec.causal: + return SupportResult.no("requires causal attention") + if spec.page_size != 1: + return SupportResult.no( + f"requires token-page KV storage (page_size=1), got {spec.page_size}" + ) + if spec.requires_attention_scores: + return SupportResult.no("does not produce per-token attention scores") + if not spec.layer_invariant_page_table: + return SupportResult.no("requires one page table shared across model layers") + if find_spec("flashinfer") is None: + return SupportResult.no("flashinfer is not installed") + try: + installed = version("flashinfer-python") + except PackageNotFoundError: + return SupportResult.no("flashinfer-python package metadata is unavailable") + numeric = tuple(int(part) for part in re.findall(r"\d+", installed)[:3]) + if numeric < (0, 6, 15): + return SupportResult.no( + f"requires flashinfer-python >= 0.6.15, got {installed}" + ) + return SupportResult.yes() + + def __init__(self) -> None: + self._state: _FlashInferPagedPrefillState | None = None + + def prepare(self, spec: PrefillAttentionOpSpec) -> None: + device = torch.device("cuda", torch.cuda.current_device()) + key = ( + device.index, + spec.num_query_heads, + spec.num_kv_heads, + spec.head_dim, + spec.activation_dtype, + ) + state = _FLASHINFER_PREFILL_STATES.get(key) + if state is None: + state = _FlashInferPagedPrefillState(device) + _FLASHINFER_PREFILL_STATES[key] = state + self._state = state + + def run( + self, + spec, + q, + view, + *, + qo_indptr, + chunk_lens, + max_context_len, + layer_idx, + ): + del chunk_lens + payload, meta = _view_parts(view) + if self._state is None: + self.prepare(spec) + state = self._state + assert state is not None + if q.dtype != spec.activation_dtype: + raise TypeError( + f"FlashInfer paged prefill expected {spec.activation_dtype} Q, got {q.dtype}." + ) + if payload.k_cache.dtype != q.dtype or payload.v_cache.dtype != q.dtype: + raise TypeError( + "FlashInfer paged prefill requires Q/K/V with the same dtype, got " + f"{q.dtype}/{payload.k_cache.dtype}/{payload.v_cache.dtype}." + ) + if meta.attn_score is not None: + raise RuntimeError( + "FlashInfer paged prefill was selected for a view that requires " + "per-token attention scores." + ) + if layer_idx == 0: + state.plan( + spec, + qo_indptr=qo_indptr, + active_slots=meta.active_slots, + req_indices=meta.req_indices, + context_lens=meta.context_lens, + max_context_len=max_context_len, + ) + elif not state.planned: + raise RuntimeError( + "FlashInfer paged prefill reached a nonzero layer before layer-0 planning." + ) + output = torch.empty_like(q) + state.wrapper.run( + q, + ( + payload.k_cache.unsqueeze(1), + payload.v_cache.unsqueeze(1), + ), + out=output, + ) + return output + + +class _FlashInferPagedPrefillState: + def __init__(self, device: torch.device) -> None: + from flashinfer.prefill import BatchPrefillWithPagedKVCacheWrapper + + self.workspace = torch.empty( + 128 * 1024 * 1024, + dtype=torch.uint8, + device=device, + ) + self.wrapper = BatchPrefillWithPagedKVCacheWrapper( + self.workspace, + kv_layout="NHD", + backend="fa3", + ) + self.planned = False + + def plan( + self, + spec: PrefillAttentionOpSpec, + *, + qo_indptr: torch.Tensor, + active_slots: torch.Tensor, + req_indices: torch.Tensor, + context_lens: torch.Tensor, + max_context_len: int, + ) -> None: + if active_slots.dim() != 2: + raise ValueError( + "FlashInfer paged prefill expects a 2D active slot table, got " + f"{tuple(active_slots.shape)}." + ) + batch_size = int(context_lens.numel()) + if batch_size <= 0 or int(req_indices.numel()) != batch_size: + raise ValueError("FlashInfer paged prefill requires matched non-empty metadata.") + max_context_len = int(max_context_len) + if max_context_len <= 0 or max_context_len > int(active_slots.shape[1]): + raise ValueError( + "FlashInfer paged prefill max context is outside the active slot table: " + f"max_context_len={max_context_len} width={int(active_slots.shape[1])}." + ) + rows = active_slots.index_select(0, req_indices.to(torch.long))[ + :, :max_context_len + ] + positions = torch.arange( + max_context_len, + device=context_lens.device, + dtype=context_lens.dtype, + ) + valid = positions.unsqueeze(0) < context_lens.unsqueeze(1) + paged_kv_indices = rows.masked_select(valid).to(torch.int32).contiguous() + zero = torch.zeros(1, device=context_lens.device, dtype=torch.int32) + paged_kv_indptr = torch.cat( + ( + zero, + context_lens.to(torch.int32).cumsum(0, dtype=torch.int32), + ) + ) + last_page_len = torch.ones( + batch_size, + device=context_lens.device, + dtype=torch.int32, + ) + self.wrapper.plan( + qo_indptr, + paged_kv_indptr, + paged_kv_indices, + last_page_len, + num_qo_heads=spec.num_query_heads, + num_kv_heads=spec.num_kv_heads, + head_dim_qk=spec.head_dim, + page_size=spec.page_size, + causal=spec.causal, + sm_scale=spec.softmax_scale, + q_data_type=spec.activation_dtype, + kv_data_type=spec.activation_dtype, + non_blocking=True, + ) + self.planned = True + + +_FLASHINFER_PREFILL_STATES: dict[tuple, _FlashInferPagedPrefillState] = {} + + +@PREFILL_ATTENTION_REGISTRY.register +class TritonPagedPrefillAttentionProvider(PrefillAttentionProvider): + name = "triton_paged_prefill" + priority = 10 + + @classmethod + def supports( + cls, spec: PrefillAttentionOpSpec, caps: DeviceCaps + ) -> SupportResult: + if caps.platform not in {PlatformEnum.CUDA, PlatformEnum.ROCM}: + return SupportResult.no(f"requires a GPU platform, got {caps.platform.name}") + if not caps.supports_triton: + return SupportResult.no("platform does not support Triton") + if spec.activation_dtype not in {torch.bfloat16, torch.float16}: + return SupportResult.no( + f"requires BF16 or FP16 Q/K/V, got {spec.activation_dtype}" + ) + if spec.head_dim not in {16, 32, 64, 128, 256}: + return SupportResult.no(f"unsupported head_dim={spec.head_dim}") + return SupportResult.yes() + + def run( + self, + spec, + q, + view, + *, + qo_indptr, + chunk_lens, + max_context_len, + layer_idx, + ): + del spec, layer_idx + payload, meta = _view_parts(view) + from sparsevllm.triton_kernel.context_flashattention_nopad import ( + context_attention_fwd, + ) + + output = torch.empty_like(q) + context_attention_fwd( + q, + payload.k_cache, + payload.v_cache, + output, + meta.req_indices, + qo_indptr[:-1], + meta.context_lens, + meta.context_lens - chunk_lens, + max_context_len, + meta.active_slots, + attn_score=meta.attn_score, + ) + return output + + +def resolve_prefill_attention_provider( + spec: PrefillAttentionOpSpec, + *, + device_index: int | None = None, +) -> PrefillAttentionProvider: + platform = platforms.current_platform + if device_index is None: + device_index = torch.cuda.current_device() if platform.is_cuda_alike() else 0 + caps = platform.get_device_caps(int(device_index)) + return OpResolver(PREFILL_ATTENTION_REGISTRY).resolve(spec, caps).provider diff --git a/src/sparsevllm/triton_kernel/gqa_flash_decoding_stage1.py b/src/sparsevllm/triton_kernel/gqa_flash_decoding_stage1.py index 140bc22d..15b3662e 100644 --- a/src/sparsevllm/triton_kernel/gqa_flash_decoding_stage1.py +++ b/src/sparsevllm/triton_kernel/gqa_flash_decoding_stage1.py @@ -328,10 +328,22 @@ def _assert_supported_layout( @torch.no_grad() def flash_decode_stage1( - q, k, v, Req_to_tokens, B_req_idx, B_Seqlen, max_len_in_batch, mid_out, mid_out_logsumexp, block_seq + q, + k, + v, + Req_to_tokens, + B_req_idx, + B_Seqlen, + max_len_in_batch, + mid_out, + mid_out_logsumexp, + block_seq, + block_n=16, + num_warps=2, + num_stages=2, ): BLOCK_SEQ = block_seq - BLOCK_N = 16 + BLOCK_N = block_n assert BLOCK_SEQ % BLOCK_N == 0 # shape constraints Lq, Lk = q.shape[-1], k.shape[-1] @@ -376,8 +388,8 @@ def flash_decode_stage1( BLOCK_SEQ=BLOCK_SEQ, BLOCK_DMODEL=Lk, BLOCK_N=BLOCK_N, - num_warps=2, - num_stages=2, + num_warps=num_warps, + num_stages=num_stages, ) return diff --git a/tests/test_operator_providers.py b/tests/test_operator_providers.py index 98e52c7f..808ad431 100644 --- a/tests/test_operator_providers.py +++ b/tests/test_operator_providers.py @@ -5,13 +5,13 @@ import pytest import torch -from sparsevllm.operators.all_reduce import HopperTp2FlashInferAllReduceProvider from sparsevllm.operators.activation import ( SILU_AND_MUL_REGISTRY, SiluAndMulSpec, TorchSiluAndMulProvider, TritonSiluAndMulProvider, ) +from sparsevllm.operators.all_reduce import ALL_REDUCE_REGISTRY, AllReduceOpSpec from sparsevllm.operators.fp8_linear import ( FP8_LINEAR_REGISTRY, FlashInferSm90Fp8LinearProvider, @@ -130,16 +130,33 @@ def _gate_up_spec(**overrides) -> GateUpSwiGLUOpSpec: return GateUpSwiGLUOpSpec(**values) -def test_flashinfer_all_reduce_falls_back_before_unsupported_shape_launch(): - provider = HopperTp2FlashInferAllReduceProvider.__new__( - HopperTp2FlashInferAllReduceProvider - ) - provider.fallback = Mock() - tensor = torch.randn(1, 248320, dtype=torch.bfloat16) - provider.fallback.run.return_value = tensor +def _all_reduce_spec(**overrides) -> AllReduceOpSpec: + values = { + "world_size": 4, + "max_tokens": 8, + "hidden_size": 3072, + "dtype": torch.bfloat16, + "cuda_graph": True, + } + values.update(overrides) + return AllReduceOpSpec(**values) + + +def test_minimax_all_reduce_prefers_flashinfer_for_profiled_h100_shape(): + with ( + patch("sparsevllm.operators.all_reduce.find_spec", return_value=object()), + patch("sparsevllm.operators.all_reduce.version", return_value="0.6.15.post1"), + ): + resolved = OpResolver(ALL_REDUCE_REGISTRY).resolve( + _all_reduce_spec(), + _cuda_caps( + (9, 0), + runtime_version="12.9", + device_name="NVIDIA H100 80GB HBM3", + ), + ) - assert provider.run(tensor) is tensor - provider.fallback.run.assert_called_once_with(tensor) + assert resolved.provider.name == "flashinfer_trtllm_sm90" @pytest.mark.parametrize("tp_size", [1, 2]) diff --git a/tests/test_parallel_context.py b/tests/test_parallel_context.py index bd256c61..30865b83 100644 --- a/tests/test_parallel_context.py +++ b/tests/test_parallel_context.py @@ -550,24 +550,6 @@ def test_dense_layers_use_tp_group_in_replicated_ep_topology(): assert embedding.weight.shape == (32, 8) -def test_parallel_group_uses_bound_all_reduce_provider(): - provider = Mock() - input_tensor, output_tensor = torch.randn(2), torch.randn(2) - provider.run.return_value = output_tensor - group = ParallelGroup( - process_group=None, - ranks=(0, 1), - rank=0, - size=2, - all_reduce_provider=provider, - ) - - actual = ParallelContext._all_reduce(input_tensor, group) - - assert actual is output_tensor - provider.run.assert_called_once_with(input_tensor) - - def test_vocab_parallel_embedding_uses_out_of_place_reduction(): reduced = torch.randn(2, 4) context = SimpleNamespace( diff --git a/tests/test_prefill_attention_provider.py b/tests/test_prefill_attention_provider.py new file mode 100644 index 00000000..cf8b63ee --- /dev/null +++ b/tests/test_prefill_attention_provider.py @@ -0,0 +1,164 @@ +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch + +from sparsevllm.operators.prefill_attention import ( + PREFILL_ATTENTION_REGISTRY, + FlashInferPagedPrefillAttentionProvider, + PrefillAttentionOpSpec, +) +from sparsevllm.operators.registry import OpResolver +from sparsevllm.platforms import DeviceCaps, PlatformEnum + + +def _spec(**overrides) -> PrefillAttentionOpSpec: + values = { + "num_query_heads": 12, + "num_kv_heads": 2, + "head_dim": 128, + "activation_dtype": torch.bfloat16, + "softmax_scale": 128**-0.5, + "causal": True, + "page_size": 1, + "requires_attention_scores": False, + "layer_invariant_page_table": True, + } + values.update(overrides) + return PrefillAttentionOpSpec(**values) + + +def _h100_caps(**overrides) -> DeviceCaps: + values = { + "platform": PlatformEnum.CUDA, + "device_type": "cuda", + "device_index": 0, + "device_name": "NVIDIA H100 80GB HBM3", + "compute_capability": (9, 0), + "runtime_version": "13.0", + "supports_graph_capture": True, + "supports_triton": True, + "supports_bfloat16": True, + "supports_native_fp8": True, + } + values.update(overrides) + return DeviceCaps(**values) + + +@patch( + "sparsevllm.operators.prefill_attention.version", + return_value="0.6.15", +) +@patch("sparsevllm.operators.prefill_attention.find_spec", return_value=object()) +def test_resolver_prefers_flashinfer_for_profiled_minimax_shape(_find, _version): + resolved = OpResolver(PREFILL_ATTENTION_REGISTRY).resolve( + _spec(), _h100_caps() + ) + assert resolved.provider.name == "flashinfer_paged_prefill_fa3_sm90" + + +@pytest.mark.parametrize( + ("spec", "caps", "reason"), + [ + (_spec(head_dim=64), _h100_caps(), "profiled local"), + (_spec(activation_dtype=torch.float16), _h100_caps(), "BF16"), + (_spec(page_size=16), _h100_caps(), "page_size=1"), + ( + _spec(requires_attention_scores=True), + _h100_caps(), + "attention scores", + ), + ( + _spec(layer_invariant_page_table=False), + _h100_caps(), + "shared across model layers", + ), + ( + _spec(), + _h100_caps(compute_capability=(8, 0), device_name="NVIDIA A100"), + "SM90", + ), + ], +) +def test_flashinfer_provider_rejects_unsupported_contracts(spec, caps, reason): + result = FlashInferPagedPrefillAttentionProvider.supports(spec, caps) + assert not result.supported + assert reason in result.reason + + +@patch( + "sparsevllm.operators.prefill_attention.version", + return_value="0.6.14", +) +@patch("sparsevllm.operators.prefill_attention.find_spec", return_value=object()) +def test_resolver_falls_back_when_flashinfer_is_too_old(_find, _version): + resolved = OpResolver(PREFILL_ATTENTION_REGISTRY).resolve( + _spec(), _h100_caps() + ) + assert resolved.provider.name == "triton_paged_prefill" + assert ( + "flashinfer_paged_prefill_fa3_sm90", + "requires flashinfer-python >= 0.6.15, got 0.6.14", + ) in resolved.rejected + + +def _torch_prefill_oracle(q, logical_k, logical_v, q_lens, kv_lens): + outputs = [] + q_cursor = 0 + for q_len, kv_len, k, v in zip(q_lens, kv_lens, logical_k, logical_v): + q_seq = q[q_cursor : q_cursor + q_len].transpose(0, 1).float() + q_cursor += q_len + k = k.transpose(0, 1).float().repeat_interleave(6, dim=0) + v = v.transpose(0, 1).float().repeat_interleave(6, dim=0) + q_positions = kv_len - q_len + torch.arange(q_len, device=q.device) + k_positions = torch.arange(kv_len, device=q.device) + allowed = k_positions.unsqueeze(0) <= q_positions.unsqueeze(1) + scores = torch.matmul(q_seq, k.transpose(-1, -2)) * (128**-0.5) + scores.masked_fill_(~allowed.unsqueeze(0), -torch.inf) + output = torch.matmul(torch.softmax(scores, dim=-1), v) + outputs.append(output.transpose(0, 1).to(torch.bfloat16)) + return torch.cat(outputs) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_flashinfer_page_size_one_matches_noncontiguous_torch_oracle(): + if torch.cuda.get_device_capability() != (9, 0): + pytest.skip("The specialized provider requires SM90.") + pytest.importorskip("flashinfer") + torch.manual_seed(20260809) + q_lens = [3, 2] + kv_lens = [5, 6] + q = torch.randn(5, 12, 128, device="cuda", dtype=torch.bfloat16) + k_cache = torch.randn(23, 2, 128, device="cuda", dtype=torch.bfloat16) + v_cache = torch.randn_like(k_cache) + pages = torch.randperm(23, device="cuda")[:11] + rows = torch.zeros(2, 6, device="cuda", dtype=torch.int32) + rows[0, :5] = pages[:5].to(torch.int32) + rows[1, :6] = pages[5:].to(torch.int32) + logical_k = [k_cache[pages[:5]], k_cache[pages[5:]]] + logical_v = [v_cache[pages[:5]], v_cache[pages[5:]]] + view = SimpleNamespace( + k_cache=k_cache, + v_cache=v_cache, + active_slots=rows, + req_indices=torch.tensor([0, 1], device="cuda", dtype=torch.int32), + context_lens=torch.tensor(kv_lens, device="cuda", dtype=torch.int32), + attn_score=None, + ) + provider = FlashInferPagedPrefillAttentionProvider() + spec = _spec() + provider.prepare(spec) + actual = provider.run( + spec, + q, + view, + qo_indptr=torch.tensor([0, 3, 5], device="cuda", dtype=torch.int32), + chunk_lens=torch.tensor(q_lens, device="cuda", dtype=torch.int32), + max_context_len=6, + layer_idx=0, + ) + expected = _torch_prefill_oracle( + q, logical_k, logical_v, q_lens, kv_lens + ) + torch.testing.assert_close(actual, expected, rtol=0.03, atol=0.03) diff --git a/tests/test_qwen35_hd256_decode_routing.py b/tests/test_qwen35_hd256_decode_routing.py index 09daeb27..642cbd06 100644 --- a/tests/test_qwen35_hd256_decode_routing.py +++ b/tests/test_qwen35_hd256_decode_routing.py @@ -17,6 +17,43 @@ class Qwen35Hd256DecodeRoutingTest(unittest.TestCase): + def test_minimax_long_decode_uses_tuned_launch_config(self): + config = TritonAttentionBackend.gqa_decode_launch_config( + block_seq=256, + max_context_len=65536, + num_heads=12, + num_kv_heads=2, + head_dim=128, + requires_attention_scores=False, + ) + + self.assertEqual(config, (1024, 128, 4)) + + def test_other_gqa_decode_shapes_keep_default_launch_config(self): + cases = ( + {"max_context_len": 32768}, + {"num_heads": 16}, + {"num_kv_heads": 4}, + {"head_dim": 256}, + {"requires_attention_scores": True}, + {"block_seq": 128}, + ) + defaults = { + "block_seq": 256, + "max_context_len": 65536, + "num_heads": 12, + "num_kv_heads": 2, + "head_dim": 128, + "requires_attention_scores": False, + } + + for override in cases: + with self.subTest(override=override): + config = TritonAttentionBackend.gqa_decode_launch_config( + **(defaults | override) + ) + self.assertEqual(config, (override.get("block_seq", 256), 16, 2)) + def _make_view(self, *, head_dim: int, attn_score=None): active_slots = torch.tensor([[0, 1, 2]], dtype=torch.int32) req_indices = torch.tensor([0], dtype=torch.int32) From 9b700000152c8ab74bd8d74d877d46b522c6e9fe Mon Sep 17 00:00:00 2001 From: CURRENTF <1206013760@qq.com> Date: Wed, 12 Aug 2026 15:03:37 +0800 Subject: [PATCH 24/48] perf: bind minimax runtime operators --- src/sparsevllm/configs/cuda_graph.py | 42 +++ src/sparsevllm/engine/decode_cuda_graph.py | 10 +- src/sparsevllm/engine/model_runner.py | 34 ++- src/sparsevllm/layers/attention.py | 42 ++- src/sparsevllm/layers/attention_backend.py | 19 -- src/sparsevllm/layers/layernorm.py | 19 ++ src/sparsevllm/models/minimax_m2.py | 271 +++++++++++++++--- src/sparsevllm/operators/all_reduce.py | 184 ++++++++++-- src/sparsevllm/operators/decode_attention.py | 158 ++++++++++ src/sparsevllm/operators/moe.py | 4 +- src/sparsevllm/operators/prefill_attention.py | 123 ++++++-- .../triton_kernel/column_parallel_rmsnorm.py | 230 +++++++++++++++ tests/test_column_parallel_rmsnorm.py | 48 ++++ tests/test_minimax_m2_model.py | 38 ++- tests/test_parallel_context.py | 22 ++ tests/test_prefill_attention_provider.py | 147 +++++++++- tests/test_prefill_schedule_policy.py | 22 ++ tests/test_qwen35_hd256_decode_routing.py | 37 --- tests/test_tp_rpc.py | 5 + 19 files changed, 1273 insertions(+), 182 deletions(-) create mode 100644 src/sparsevllm/operators/decode_attention.py create mode 100644 src/sparsevllm/triton_kernel/column_parallel_rmsnorm.py diff --git a/src/sparsevllm/configs/cuda_graph.py b/src/sparsevllm/configs/cuda_graph.py index 537e0c61..3e99aa92 100644 --- a/src/sparsevllm/configs/cuda_graph.py +++ b/src/sparsevllm/configs/cuda_graph.py @@ -81,6 +81,48 @@ def _resolve_decode_cuda_graph_capture_sizes( return sizes +def _select_decode_cuda_graph_batch_size( + real_batch_size: int, + capture_sizes: list[int] | tuple[int, ...], +) -> int: + real_batch_size = int(real_batch_size) + if real_batch_size <= 0: + raise ValueError( + f"decode batch size must be > 0, got {real_batch_size}." + ) + sizes = sorted(set(int(size) for size in capture_sizes)) + if not sizes or any(size <= 0 for size in sizes): + raise ValueError( + "decode_cuda_graph_capture_sizes must contain positive integers, " + f"got {sizes}." + ) + for size in sizes: + if size >= real_batch_size: + return size + raise ValueError( + "decode_cuda_graph capture sizes do not cover current decode batch: " + f"batch_size={real_batch_size}, capture_sizes={sizes}." + ) + + +def _resolve_decode_static_batch_capacity( + capture_sizes: list[int] | tuple[int, ...], + *, + max_num_seqs_in_batch: int, + max_decoding_seqs: int, +) -> int: + """Return the largest padded decode batch reachable by the scheduler.""" + + max_real_batch_size = min( + int(max_num_seqs_in_batch), + int(max_decoding_seqs), + ) + return _select_decode_cuda_graph_batch_size( + max_real_batch_size, + capture_sizes, + ) + + def _default_decode_cuda_graph_context_sizes(max_model_len: int) -> list[int]: """Default decode graph context buckets: 1k, 2k, 4k, ... up to max_model_len.""" max_model_len = int(max_model_len) diff --git a/src/sparsevllm/engine/decode_cuda_graph.py b/src/sparsevllm/engine/decode_cuda_graph.py index f0b6a963..e94a7344 100644 --- a/src/sparsevllm/engine/decode_cuda_graph.py +++ b/src/sparsevllm/engine/decode_cuda_graph.py @@ -7,6 +7,7 @@ import torch from sparsevllm.engine.sequence import Sequence +from sparsevllm.configs.cuda_graph import _select_decode_cuda_graph_batch_size import sparsevllm.platforms as platforms from sparsevllm.utils.context import get_context, set_context from sparsevllm.utils.profiler import profiler @@ -212,12 +213,9 @@ def _select_graph_batch_size(self, real_batch_size: int) -> int: if selected is not None: return int(selected) - for size in self.capture_sizes: - if size >= real_batch_size: - return int(size) - raise ValueError( - "decode_cuda_graph capture sizes do not cover current decode batch: " - f"batch_size={real_batch_size}, capture_sizes={self.capture_sizes}." + return _select_decode_cuda_graph_batch_size( + real_batch_size, + self.capture_sizes, ) def _select_state( diff --git a/src/sparsevllm/engine/model_runner.py b/src/sparsevllm/engine/model_runner.py index bd6bd211..e77b2829 100644 --- a/src/sparsevllm/engine/model_runner.py +++ b/src/sparsevllm/engine/model_runner.py @@ -13,6 +13,7 @@ _resolve_decode_cuda_graph_capture_sizes, _resolve_decode_cuda_graph_context_sizes, ) +from sparsevllm.configs.cuda_graph import _resolve_decode_static_batch_capacity from sparsevllm.distributed import init_parallel_context, reset_parallel_context from sparsevllm.engine.sequence import Sequence from sparsevllm.models.qwen2 import Qwen2ForCausalLM @@ -50,9 +51,13 @@ Glm4MoeLiteForCausalLM = None try: - from sparsevllm.models.minimax_m2 import MiniMaxM2ForCausalLM + from sparsevllm.models.minimax_m2 import ( + MiniMaxM2ForCausalLM, + build_minimax_m2_runtime_config, + ) except ImportError: MiniMaxM2ForCausalLM = None + build_minimax_m2_runtime_config = None try: from sparsevllm.models.qwen3_5 import Qwen35ForCausalLM @@ -181,6 +186,10 @@ def __init__( "decode_cuda_graph", bool(getattr(config, "decode_cuda_graph", False)), ) + decode_static_capture_sizes = _resolve_decode_cuda_graph_capture_sizes( + config.decode_cuda_graph_capture_sizes, + config.max_decoding_seqs, + ) model_kwargs = {} if config.model_spec.runtime_class_name == "Glm4MoeLiteForCausalLM": from sparsevllm.models.glm4_moe_lite import ( @@ -203,6 +212,21 @@ def __init__( "decode_cuda_graph": config.decode_cuda_graph, "expect_mtp_weights": not config.tiny_random, } + elif config.model_spec.runtime_class_name == "MiniMaxM2ForCausalLM": + if build_minimax_m2_runtime_config is None: + raise ImportError("MiniMax M2 runtime operators are unavailable.") + model_kwargs["runtime_config"] = build_minimax_m2_runtime_config( + hf_config, + self.parallel_context, + layer_invariant_page_table=config.vllm_sparse_method == "", + max_decode_tokens=_resolve_decode_static_batch_capacity( + decode_static_capture_sizes, + max_num_seqs_in_batch=config.max_num_seqs_in_batch, + max_decoding_seqs=config.max_decoding_seqs, + ), + cuda_graph=config.decode_cuda_graph, + device_index=int(self.device.index or 0), + ) self.model = _create_model(hf_config, config.model_spec, **model_kwargs) if config.tiny_random: @@ -304,10 +328,6 @@ def __init__( # 加载 DeltaKV 压缩器 self.load_deltakv_compressors() - decode_static_capture_sizes = _resolve_decode_cuda_graph_capture_sizes( - self.config.decode_cuda_graph_capture_sizes, - self.config.max_decoding_seqs, - ) decode_static_context_sizes = _resolve_decode_cuda_graph_context_sizes( self.config.decode_cuda_graph_context_sizes, self.config.max_model_len, @@ -356,6 +376,10 @@ def exit(self): if self.config.decode_cuda_graph: self.decode_cuda_graph_runner.clear_captured_graphs() self.platform.synchronize() + close_runtime_operators = getattr(self.model, "close_runtime_operators", None) + if callable(close_runtime_operators): + close_runtime_operators() + self.platform.synchronize() if self.world_size > 1: self.shm.close() self.parallel_context.world_barrier( diff --git a/src/sparsevllm/layers/attention.py b/src/sparsevllm/layers/attention.py index 6f84e48a..bce1d893 100644 --- a/src/sparsevllm/layers/attention.py +++ b/src/sparsevllm/layers/attention.py @@ -5,10 +5,8 @@ from sparsevllm.engine.cache_manager import ExplicitKVPayload from sparsevllm.layers.attention_backend import TritonAttentionBackend -from sparsevllm.operators.prefill_attention import ( - PrefillAttentionOpSpec, - resolve_prefill_attention_provider, -) +from sparsevllm.operators.decode_attention import PreparedDecodeAttentionLaunchOp +from sparsevllm.operators.prefill_attention import PreparedPrefillAttentionOp from sparsevllm.utils.context import get_context from sparsevllm.engine.sparse_controller import SparseController @@ -62,7 +60,8 @@ def __init__( scale, num_kv_heads, *, - prefill_op_spec: PrefillAttentionOpSpec | None = None, + prefill_op: PreparedPrefillAttentionOp | None = None, + decode_launch_op: PreparedDecodeAttentionLaunchOp | None = None, ): super().__init__() self.num_heads = num_heads @@ -70,14 +69,8 @@ def __init__( self.scale = scale self.num_kv_heads = num_kv_heads self.attention_backend = TritonAttentionBackend() - self.prefill_op_spec = prefill_op_spec - self.prefill_provider = ( - resolve_prefill_attention_provider(prefill_op_spec) - if prefill_op_spec is not None - else None - ) - if self.prefill_provider is not None: - self.prefill_provider.prepare(prefill_op_spec) + self.prefill_op = prefill_op + self.decode_launch_op = decode_launch_op def forward( self, @@ -130,7 +123,7 @@ def forward( ) if fake_output is not None: o = fake_output - elif self.prefill_provider is None: + elif self.prefill_op is None: o = self.attention_backend.run_prefill( q, prefill_view, @@ -144,8 +137,7 @@ def forward( prefill_view, chunk_lens=chunk_lens, ) - o = self.prefill_provider.run( - self.prefill_op_spec, + o = self.prefill_op.run( q, prefill_view, qo_indptr=context.cu_seqlens_q, @@ -222,16 +214,16 @@ def forward( f"decode requires a positive context length, got {max_len_in_batch} at layer={layer_idx}" ) BLOCK_SEQ = cache_manager.get_decode_block_seq(layer_idx, 256) - BLOCK_SEQ, gqa_block_n, gqa_num_warps = ( - self.attention_backend.gqa_decode_launch_config( - block_seq=BLOCK_SEQ, - max_context_len=max_len_in_batch, - num_heads=self.num_heads, - num_kv_heads=self.num_kv_heads, - head_dim=self.head_dim, - requires_attention_scores=decode_meta.attn_score is not None, + if self.decode_launch_op is None: + gqa_block_n, gqa_num_warps = 16, 2 + else: + BLOCK_SEQ, gqa_block_n, gqa_num_warps = ( + self.decode_launch_op.launch_config( + block_seq=BLOCK_SEQ, + max_context_len=max_len_in_batch, + requires_attention_scores=decode_meta.attn_score is not None, + ) ) - ) num_seq_blocks = (max_len_in_batch + BLOCK_SEQ - 1) // BLOCK_SEQ mid_o, mid_o_logexpsum = get_decode_workspace( diff --git a/src/sparsevllm/layers/attention_backend.py b/src/sparsevllm/layers/attention_backend.py index cb3ecd3e..5d13e474 100644 --- a/src/sparsevllm/layers/attention_backend.py +++ b/src/sparsevllm/layers/attention_backend.py @@ -118,25 +118,6 @@ class TritonAttentionBackend: def __init__(self) -> None: record_operator_binding("Attention", self) - @staticmethod - def gqa_decode_launch_config( - *, - block_seq: int, - max_context_len: int, - num_heads: int, - num_kv_heads: int, - head_dim: int, - requires_attention_scores: bool, - ) -> tuple[int, int, int]: - if ( - block_seq == 256 - and max_context_len > 32768 - and (num_heads, num_kv_heads, head_dim) == (12, 2, 128) - and not requires_attention_scores - ): - return 1024, 128, 4 - return block_seq, 16, 2 - def maybe_run_fake_prefill( self, q: torch.Tensor, diff --git a/src/sparsevllm/layers/layernorm.py b/src/sparsevllm/layers/layernorm.py index c4efbc25..a7afb418 100755 --- a/src/sparsevllm/layers/layernorm.py +++ b/src/sparsevllm/layers/layernorm.py @@ -222,6 +222,25 @@ def forward_pair( ) -> tuple[torch.Tensor, torch.Tensor]: if self.parallel_context is not other_norm.parallel_context: raise ValueError("Paired column-parallel RMSNorms must share a context.") + if x.is_cuda or other.is_cuda: + from sparsevllm.triton_kernel.column_parallel_rmsnorm import ( + paired_rms_apply, + paired_square_sums, + ) + + square_sums = paired_square_sums(x, other) + self.parallel_context.attention_tp_all_reduce(square_sums) + return paired_rms_apply( + x, + other, + square_sums, + self.weight, + other_norm.weight, + x_global_hidden_size=self.global_hidden_size, + other_global_hidden_size=other_norm.global_hidden_size, + x_eps=self.eps, + other_eps=other_norm.eps, + ) square_sums = torch.stack( ( x.float().square().sum(dim=-1), diff --git a/src/sparsevllm/models/minimax_m2.py b/src/sparsevllm/models/minimax_m2.py index b07a8a1c..a7e2383f 100644 --- a/src/sparsevllm/models/minimax_m2.py +++ b/src/sparsevllm/models/minimax_m2.py @@ -1,12 +1,18 @@ from __future__ import annotations import re +from dataclasses import dataclass +from functools import partial import torch import torch.nn.functional as F from torch import nn -from sparsevllm.distributed import get_parallel_context +from sparsevllm.distributed import ( + ParallelContext, + ParallelGroup, + get_parallel_context, +) from sparsevllm.layers.attention import Attention from sparsevllm.layers.embed_head import ParallelLMHead from sparsevllm.layers.expert_weights import PackedExpertWeightLoader @@ -22,7 +28,21 @@ model_activation_dtype, resolve_moe_provider, ) -from sparsevllm.operators.prefill_attention import PrefillAttentionOpSpec +from sparsevllm.operators.all_reduce import ( + AllReduceOpSpec, + PreparedAllReduceOp, + prepare_all_reduce_op, +) +from sparsevllm.operators.decode_attention import ( + DecodeAttentionLaunchSpec, + PreparedDecodeAttentionLaunchOp, + prepare_decode_attention_launch_op, +) +from sparsevllm.operators.prefill_attention import ( + PrefillAttentionOpSpec, + PreparedPrefillAttentionOp, + prepare_prefill_attention_op, +) from sparsevllm.platforms import device_runtime from sparsevllm.quantization.fp8_tp import Fp8ExpertTpShard from sparsevllm.utils.context import get_context @@ -40,6 +60,127 @@ ) +@dataclass +class MiniMaxM2RuntimeConfig: + prefill_attention_op: PreparedPrefillAttentionOp + decode_launch_op: PreparedDecodeAttentionLaunchOp + attention_decode_all_reduce: PreparedAllReduceOp + moe_decode_all_reduce: PreparedAllReduceOp + cuda_graph: bool + _closed: bool = False + + def close(self) -> None: + if self._closed: + return + self.prefill_attention_op.close() + seen: set[int] = set() + for op in ( + self.attention_decode_all_reduce, + self.moe_decode_all_reduce, + ): + if id(op) in seen: + continue + seen.add(id(op)) + op.close() + self._closed = True + + +def _prepare_decode_all_reduce( + group: ParallelGroup, + *, + max_tokens: int, + hidden_size: int, + dtype: torch.dtype, + cuda_graph: bool, + device_index: int, +) -> PreparedAllReduceOp: + return prepare_all_reduce_op( + AllReduceOpSpec( + world_size=group.size, + max_tokens=max_tokens, + hidden_size=hidden_size, + dtype=dtype, + cuda_graph=cuda_graph, + ), + group=group.process_group, + rank=group.rank, + device_index=device_index, + ) + + +def build_minimax_m2_runtime_config( + config, + parallel_context: ParallelContext, + *, + layer_invariant_page_table: bool, + max_decode_tokens: int, + cuda_graph: bool, + device_index: int, +) -> MiniMaxM2RuntimeConfig: + tp_size = int(parallel_context.attention_tp_size) + if ( + int(config.num_attention_heads) % tp_size + or int(config.num_key_value_heads) % tp_size + ): + raise ValueError( + "MiniMax attention heads must be divisible by attention TP size." + ) + num_query_heads = int(config.num_attention_heads) // tp_size + num_kv_heads = int(config.num_key_value_heads) // tp_size + head_dim = int(config.head_dim) + activation_dtype = model_activation_dtype(config) + prefill_attention_op = prepare_prefill_attention_op( + PrefillAttentionOpSpec( + num_query_heads=num_query_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + activation_dtype=activation_dtype, + softmax_scale=head_dim**-0.5, + causal=True, + page_size=1, + requires_attention_scores=False, + layer_invariant_page_table=bool(layer_invariant_page_table), + ), + device_index=device_index, + ) + decode_launch_op = prepare_decode_attention_launch_op( + DecodeAttentionLaunchSpec( + num_query_heads=num_query_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + activation_dtype=activation_dtype, + page_size=1, + ), + device_index=device_index, + ) + moe_decode_all_reduce = _prepare_decode_all_reduce( + parallel_context.world, + max_tokens=int(max_decode_tokens), + hidden_size=int(config.hidden_size), + dtype=activation_dtype, + cuda_graph=bool(cuda_graph), + device_index=device_index, + ) + if parallel_context.attention.ranks == parallel_context.world.ranks: + attention_decode_all_reduce = moe_decode_all_reduce + else: + attention_decode_all_reduce = _prepare_decode_all_reduce( + parallel_context.attention, + max_tokens=int(max_decode_tokens), + hidden_size=int(config.hidden_size), + dtype=activation_dtype, + cuda_graph=bool(cuda_graph), + device_index=device_index, + ) + return MiniMaxM2RuntimeConfig( + prefill_attention_op=prefill_attention_op, + decode_launch_op=decode_launch_op, + attention_decode_all_reduce=attention_decode_all_reduce, + moe_decode_all_reduce=moe_decode_all_reduce, + cuda_graph=bool(cuda_graph), + ) + + class MiniMaxM2Router(nn.Module): def __init__(self, config) -> None: super().__init__() @@ -74,7 +215,12 @@ def forward( class MiniMaxM2PackedExperts(PackedExpertWeightLoader, nn.Module): checkpoint_projection_map = {"w1": "gate", "w2": "down", "w3": "up"} - def __init__(self, config) -> None: + def __init__( + self, + config, + *, + cuda_graph: bool | None = None, + ) -> None: super().__init__() parallel_context = get_parallel_context() self.tp_rank = int(parallel_context.moe_tp_rank) @@ -114,7 +260,11 @@ def __init__(self, config) -> None: weight_dtype=torch.float8_e4m3fn, block_shape=tuple(config.quantization_config.weight_block_size), ep_size=self.ep_size, - cuda_graph=bool(getattr(config, "decode_cuda_graph", False)), + cuda_graph=( + bool(getattr(config, "decode_cuda_graph", False)) + if cuda_graph is None + else bool(cuda_graph) + ), tp_size=self.tp_size, routing_method="biased_sigmoid", scale_dtype=torch.float32, @@ -253,16 +403,24 @@ def forward( class MiniMaxM2SparseMoeBlock(nn.Module): - def __init__(self, config) -> None: + def __init__( + self, + config, + runtime_config: MiniMaxM2RuntimeConfig | None = None, + ) -> None: super().__init__() self.parallel_context = get_parallel_context() + self.runtime_config = runtime_config self.mlp_chunk_size = int(getattr(config, "mlp_chunk_size", 16384)) if self.mlp_chunk_size <= 0: raise ValueError( f"mlp_chunk_size must be > 0, got {self.mlp_chunk_size}." ) self.gate = MiniMaxM2Router(config) - self.experts = MiniMaxM2PackedExperts(config) + self.experts = MiniMaxM2PackedExperts( + config, + cuda_graph=(None if runtime_config is None else runtime_config.cuda_graph), + ) @property def e_score_correction_bias(self) -> nn.Parameter: @@ -288,13 +446,22 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: self.experts(chunk, topk_ids, topk_weights) ) local_output = torch.cat(local_output_chunks, dim=0) + if self.runtime_config is not None: + context = get_context() + if not context.is_prefill: + return self.runtime_config.moe_decode_all_reduce.run(local_output) return self.parallel_context.world_all_reduce(local_output) class MiniMaxM2Attention(nn.Module): - def __init__(self, config) -> None: + def __init__( + self, + config, + runtime_config: MiniMaxM2RuntimeConfig | None = None, + ) -> None: super().__init__() self.parallel_context = get_parallel_context() + self.runtime_config = runtime_config tp_size = int(self.parallel_context.tp_size) self.total_num_heads = int(config.num_attention_heads) self.total_num_kv_heads = int(config.num_key_value_heads) @@ -319,6 +486,7 @@ def __init__(self, config) -> None: int(config.hidden_size), bias=False, quantization=config.quantization_config, + reduce_results=runtime_config is None, ) self.q_norm = ColumnParallelRMSNorm( self.total_num_heads * self.head_dim, @@ -337,27 +505,17 @@ def __init__(self, config) -> None: base=float(config.rope_theta), rope_scaling=None, ) - prefill_op_spec = None - if hasattr(config, "prefill_kv_view_layer_invariant"): - prefill_op_spec = PrefillAttentionOpSpec( - num_query_heads=self.num_heads, - num_kv_heads=self.num_kv_heads, - head_dim=self.head_dim, - activation_dtype=model_activation_dtype(config), - softmax_scale=self.head_dim**-0.5, - causal=True, - page_size=1, - requires_attention_scores=False, - layer_invariant_page_table=bool( - config.prefill_kv_view_layer_invariant - ), - ) self.attn = Attention( self.num_heads, self.head_dim, self.head_dim**-0.5, self.num_kv_heads, - prefill_op_spec=prefill_op_spec, + prefill_op=( + None if runtime_config is None else runtime_config.prefill_attention_op + ), + decode_launch_op=( + None if runtime_config is None else runtime_config.decode_launch_op + ), ) def forward( @@ -382,15 +540,24 @@ def forward( ) context.cache_manager.save_rope_kv_if_needed(layer_idx, k, v) output = self.attn(q, k, v).flatten(1, -1) - return self.o_proj(output) + output = self.o_proj(output) + if self.runtime_config is not None: + if context.is_prefill: + return self.parallel_context.attention_tp_all_reduce(output) + return self.runtime_config.attention_decode_all_reduce.run(output) + return output class MiniMaxM2DecoderLayer(nn.Module): - def __init__(self, config) -> None: + def __init__( + self, + config, + runtime_config: MiniMaxM2RuntimeConfig | None = None, + ) -> None: super().__init__() self.parallel_context = get_parallel_context() - self.self_attn = MiniMaxM2Attention(config) - self.block_sparse_moe = MiniMaxM2SparseMoeBlock(config) + self.self_attn = MiniMaxM2Attention(config, runtime_config) + self.block_sparse_moe = MiniMaxM2SparseMoeBlock(config, runtime_config) self.input_layernorm = RMSNorm( config.hidden_size, eps=config.rms_norm_eps, @@ -422,8 +589,16 @@ def forward( class MiniMaxM2Model(Qwen3ModelBase): - def __init__(self, config) -> None: - super().__init__(config, MiniMaxM2DecoderLayer) + def __init__( + self, + config, + runtime_config: MiniMaxM2RuntimeConfig | None = None, + ) -> None: + layer_factory = partial( + MiniMaxM2DecoderLayer, + runtime_config=runtime_config, + ) + super().__init__(config, layer_factory) self.norm = RMSNorm( config.hidden_size, eps=config.rms_norm_eps, @@ -438,15 +613,24 @@ class MiniMaxM2ForCausalLM(nn.Module): "v_proj": ("qkv_proj", "v"), } - def __init__(self, config) -> None: + def __init__( + self, + config, + runtime_config: MiniMaxM2RuntimeConfig | None = None, + ) -> None: super().__init__() self.config = config self.parallel_context = get_parallel_context() - self.model = MiniMaxM2Model(config) + self.runtime_config = runtime_config + self.model = MiniMaxM2Model(config, runtime_config) self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size) self._intentionally_skipped_expert_weights: set[str] = set() self._intentionally_skipped_expert_scales: set[str] = set() + def close_runtime_operators(self) -> None: + if self.runtime_config is not None: + self.runtime_config.close() + @torch.inference_mode() def warmup_moe(self, num_tokens: int = 1) -> None: num_tokens = int(num_tokens) @@ -657,19 +841,30 @@ def validate_loaded_weights(self, loaded_parameter_names: set[str]) -> None: "Unexpectedly skipped MiniMax expert entries: " f"weights={unexpected_skips[:4]}, scales={unexpected_scale_skips[:4]}." ) + prefill_provider = ( + self.runtime_config.prefill_attention_op.name + if self.runtime_config is not None + else "legacy_triton" + ) + if self.runtime_config is None: + all_reduce_providers = "legacy_torch_distributed" + else: + all_reduce_providers = ( + "attention=" + f"{self.runtime_config.attention_decode_all_reduce.name}," + "moe=" + f"{self.runtime_config.moe_decode_all_reduce.name}" + ) logger.info( - "Loaded MiniMax M2 rank {} provider={} prefill_provider={} all_reduce_provider={} " + "Loaded MiniMax M2 rank {} provider={} prefill_provider={} " + "all_reduce_providers={} " "attention TP {}/{} MoE TP " "{}/{} local experts [{}, {}) across {} layers; intentionally skipped " "{} remote expert weight/scale pairs.", self.parallel_context.world_rank, self.model.layers[0].block_sparse_moe.experts.provider.name, - getattr( - self.model.layers[0].self_attn.attn.prefill_provider, - "name", - "legacy_triton", - ), - getattr(self.parallel_context.all_reduce_provider, "name", "unconfigured"), + prefill_provider, + all_reduce_providers, self.parallel_context.tp_rank, self.parallel_context.tp_size, self.parallel_context.moe_tp_rank, diff --git a/src/sparsevllm/operators/all_reduce.py b/src/sparsevllm/operators/all_reduce.py index 153fce91..e359ce3f 100644 --- a/src/sparsevllm/operators/all_reduce.py +++ b/src/sparsevllm/operators/all_reduce.py @@ -39,22 +39,23 @@ def prepare( self, spec: AllReduceOpSpec, *, - group: dist.ProcessGroup, + group: dist.ProcessGroup | None, rank: int, + device_index: int | None = None, ) -> None: - del spec, group, rank + del spec, group, rank, device_index - def can_run(self, spec: AllReduceOpSpec, tensor: torch.Tensor) -> bool: - del spec, tensor - return True + def close(self) -> None: + pass def run( self, spec: AllReduceOpSpec, tensor: torch.Tensor, *, - group: dist.ProcessGroup, + group: dist.ProcessGroup | None, ) -> torch.Tensor: + """Return the reduced tensor, which may or may not alias the input.""" raise NotImplementedError @@ -114,11 +115,24 @@ def supports(cls, spec: AllReduceOpSpec, caps: DeviceCaps) -> SupportResult: def __init__(self) -> None: self.workspace = None + self._output_buffer: torch.Tensor | None = None - def prepare(self, spec, *, group, rank) -> None: + def prepare(self, spec, *, group, rank, device_index=None) -> None: from flashinfer.comm import create_allreduce_fusion_workspace - self.workspace = create_allreduce_fusion_workspace( + if self.workspace is not None or self._output_buffer is not None: + raise RuntimeError("FlashInfer all-reduce provider is already prepared.") + if group is None: + raise RuntimeError("FlashInfer all-reduce requires a distributed process group.") + current_device = torch.cuda.current_device() + if device_index is None: + device_index = current_device + if int(device_index) != current_device: + raise RuntimeError( + "FlashInfer all-reduce must be prepared on the selected CUDA device: " + f"selected={device_index} current={current_device}." + ) + workspace = create_allreduce_fusion_workspace( backend="trtllm", world_size=spec.world_size, rank=rank, @@ -127,34 +141,44 @@ def prepare(self, spec, *, group, rank) -> None: dtype=spec.dtype, group=group, ) + try: + output_buffer = torch.empty( + (spec.max_tokens, spec.hidden_size), + dtype=spec.dtype, + device=torch.device("cuda", int(device_index)), + ) + except Exception: + workspace.destroy() + raise + self.workspace = workspace + self._output_buffer = output_buffer - def can_run(self, spec, tensor) -> bool: - return ( - tensor.is_cuda - and tensor.dtype == spec.dtype - and tensor.ndim == 2 - and 0 < int(tensor.shape[0]) <= spec.max_tokens - and int(tensor.shape[1]) == spec.hidden_size - and tensor.is_contiguous() - ) + def close(self) -> None: + workspace = self.workspace + self.workspace = None + self._output_buffer = None + if workspace is not None: + workspace.destroy() def run(self, spec, tensor, *, group) -> torch.Tensor: del group - if self.workspace is None: + if self.workspace is None or self._output_buffer is None: raise RuntimeError("FlashInfer all-reduce provider was not prepared.") - if not self.can_run(spec, tensor): + if not tensor.is_cuda: raise ValueError( - "FlashInfer all-reduce received an unsupported tensor: " - f"shape={tuple(tensor.shape)} dtype={tensor.dtype} device={tensor.device}." + f"FlashInfer all-reduce requires a CUDA tensor, got {tensor.device}." ) from flashinfer.comm import allreduce_fusion from flashinfer.comm.trtllm_ar import AllReduceFusionPattern - return allreduce_fusion( + output = self._output_buffer[: int(tensor.shape[0])] + allreduce_fusion( input=tensor, workspace=self.workspace, pattern=AllReduceFusionPattern.kAllReduce, + output=output, ) + return output @ALL_REDUCE_REGISTRY.register @@ -168,13 +192,121 @@ def supports(cls, spec: AllReduceOpSpec, caps: DeviceCaps) -> SupportResult: return SupportResult.yes() def run(self, spec, tensor, *, group) -> torch.Tensor: - del spec - dist.all_reduce(tensor, group=group) + if spec.world_size > 1: + if group is None: + raise RuntimeError( + "Torch distributed all-reduce requires a process group when world_size > 1." + ) + dist.all_reduce(tensor, group=group) return tensor -def resolve_all_reduce_provider(spec: AllReduceOpSpec) -> AllReduceProvider: +def _validate_tensor_contract(spec: AllReduceOpSpec, tensor: torch.Tensor) -> None: + if tensor.dtype != spec.dtype: + raise TypeError( + f"All-reduce expected dtype={spec.dtype}, got {tensor.dtype}." + ) + if tensor.ndim != 2: + raise ValueError( + f"All-reduce expects [tokens, hidden], got shape={tuple(tensor.shape)}." + ) + token_count, hidden_size = (int(value) for value in tensor.shape) + if not 0 < token_count <= spec.max_tokens: + raise ValueError( + "All-reduce token count is outside the prepared range: " + f"tokens={token_count} max_tokens={spec.max_tokens}." + ) + if hidden_size != spec.hidden_size: + raise ValueError( + "All-reduce hidden size does not match the prepared operator: " + f"hidden={hidden_size} expected={spec.hidden_size}." + ) + if not tensor.is_contiguous(): + raise ValueError( + f"All-reduce requires a contiguous tensor, got stride={tensor.stride()}." + ) + + +class PreparedAllReduceOp: + """A pre-bound all-reduce with no execution-time fallback.""" + + def __init__( + self, + spec: AllReduceOpSpec, + provider: AllReduceProvider, + *, + group: dist.ProcessGroup | None, + ) -> None: + self.spec = spec + self.provider = provider + self.group = group + self._closed = False + + @property + def name(self) -> str: + return self.provider.name + + def run(self, tensor: torch.Tensor) -> torch.Tensor: + if self._closed: + raise RuntimeError("All-reduce operator is closed.") + _validate_tensor_contract(self.spec, tensor) + output = self.provider.run(self.spec, tensor, group=self.group) + if not isinstance(output, torch.Tensor): + raise TypeError( + f"All-reduce provider {self.provider.name} returned " + f"{type(output).__name__}, expected torch.Tensor." + ) + _validate_tensor_contract(self.spec, output) + if output.shape != tensor.shape or output.device != tensor.device: + raise ValueError( + f"All-reduce provider {self.provider.name} returned an incompatible tensor: " + f"input_shape={tuple(tensor.shape)} output_shape={tuple(output.shape)} " + f"input_device={tensor.device} output_device={output.device}." + ) + return output + + def close(self) -> None: + if self._closed: + return + self.provider.close() + self._closed = True + + +def resolve_all_reduce_provider( + spec: AllReduceOpSpec, + *, + device_index: int | None = None, +) -> AllReduceProvider: platform = platforms.current_platform - device_index = torch.cuda.current_device() if platform.is_cuda_alike() else 0 + if device_index is None: + device_index = torch.cuda.current_device() if platform.is_cuda_alike() else 0 caps = platform.get_device_caps(int(device_index)) return OpResolver(ALL_REDUCE_REGISTRY).resolve(spec, caps).provider + + +def prepare_all_reduce_op( + spec: AllReduceOpSpec, + *, + group: dist.ProcessGroup | None, + rank: int, + device_index: int | None = None, + provider: AllReduceProvider | None = None, +) -> PreparedAllReduceOp: + rank = int(rank) + if not 0 <= rank < spec.world_size: + raise ValueError( + f"All-reduce rank must be in [0, {spec.world_size}), got {rank}." + ) + if spec.world_size > 1 and group is None: + raise RuntimeError( + "All-reduce requires a process group when world_size > 1." + ) + if provider is None: + provider = resolve_all_reduce_provider(spec, device_index=device_index) + provider.prepare( + spec, + group=group, + rank=rank, + device_index=device_index, + ) + return PreparedAllReduceOp(spec, provider, group=group) diff --git a/src/sparsevllm/operators/decode_attention.py b/src/sparsevllm/operators/decode_attention.py new file mode 100644 index 00000000..f94c1064 --- /dev/null +++ b/src/sparsevllm/operators/decode_attention.py @@ -0,0 +1,158 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +import sparsevllm.platforms as platforms +from sparsevllm.operators.registry import OpRegistry, OpResolver, SupportResult +from sparsevllm.platforms.interface import DeviceCaps, PlatformEnum + + +@dataclass(frozen=True) +class DecodeAttentionLaunchSpec: + num_query_heads: int + num_kv_heads: int + head_dim: int + activation_dtype: torch.dtype + page_size: int = 1 + + def __post_init__(self) -> None: + if self.num_query_heads <= 0 or self.num_kv_heads <= 0: + raise ValueError("Decode attention head counts must be positive.") + if self.num_query_heads % self.num_kv_heads: + raise ValueError("Decode query heads must be divisible by KV heads.") + if self.head_dim <= 0 or self.page_size <= 0: + raise ValueError("Decode attention dimensions must be positive.") + + +class DecodeAttentionLaunchProvider: + name = "" + priority = 0 + + def launch_config( + self, + *, + block_seq: int, + max_context_len: int, + requires_attention_scores: bool, + ) -> tuple[int, int, int]: + raise NotImplementedError + + +DECODE_ATTENTION_LAUNCH_REGISTRY: OpRegistry[ + DecodeAttentionLaunchSpec, DecodeAttentionLaunchProvider +] = OpRegistry("decode attention launch") + + +@DECODE_ATTENTION_LAUNCH_REGISTRY.register +class H100LongGqaDecodeLaunchProvider(DecodeAttentionLaunchProvider): + name = "h100_long_gqa_12q_2kv_hd128" + priority = 100 + + @classmethod + def supports( + cls, + spec: DecodeAttentionLaunchSpec, + caps: DeviceCaps, + ) -> SupportResult: + if caps.platform != PlatformEnum.CUDA or caps.compute_capability != (9, 0): + return SupportResult.no( + f"requires CUDA SM90, got {caps.platform.name} {caps.compute_capability}" + ) + if caps.device_name != "NVIDIA H100 80GB HBM3": + return SupportResult.no( + "requires profiled NVIDIA H100 80GB HBM3 hardware, " + f"got {caps.device_name}" + ) + if not caps.supports_triton: + return SupportResult.no("platform does not support Triton") + if spec.activation_dtype != torch.bfloat16: + return SupportResult.no( + f"requires BF16 query/KV tensors, got {spec.activation_dtype}" + ) + expected_shape = (12, 2, 128) + actual_shape = ( + spec.num_query_heads, + spec.num_kv_heads, + spec.head_dim, + ) + if actual_shape != expected_shape: + return SupportResult.no( + f"requires profiled local Q/KV/head shape {expected_shape}, got {actual_shape}" + ) + if spec.page_size != 1: + return SupportResult.no( + f"requires token-page KV storage (page_size=1), got {spec.page_size}" + ) + return SupportResult.yes() + + def launch_config( + self, + *, + block_seq: int, + max_context_len: int, + requires_attention_scores: bool, + ) -> tuple[int, int, int]: + if ( + int(block_seq) == 256 + and int(max_context_len) > 32768 + and not requires_attention_scores + ): + return 1024, 128, 4 + return int(block_seq), 16, 2 + + +@DECODE_ATTENTION_LAUNCH_REGISTRY.register +class DefaultGqaDecodeLaunchProvider(DecodeAttentionLaunchProvider): + name = "default_gqa" + priority = 10 + + @classmethod + def supports( + cls, + spec: DecodeAttentionLaunchSpec, + caps: DeviceCaps, + ) -> SupportResult: + del spec, caps + return SupportResult.yes() + + def launch_config( + self, + *, + block_seq: int, + max_context_len: int, + requires_attention_scores: bool, + ) -> tuple[int, int, int]: + del max_context_len, requires_attention_scores + return int(block_seq), 16, 2 + + +class PreparedDecodeAttentionLaunchOp: + def __init__( + self, + spec: DecodeAttentionLaunchSpec, + provider: DecodeAttentionLaunchProvider, + ) -> None: + self.spec = spec + self.provider = provider + + @property + def name(self) -> str: + return self.provider.name + + def launch_config(self, **kwargs) -> tuple[int, int, int]: + return self.provider.launch_config(**kwargs) + + +def prepare_decode_attention_launch_op( + spec: DecodeAttentionLaunchSpec, + *, + device_index: int | None = None, +) -> PreparedDecodeAttentionLaunchOp: + platform = platforms.current_platform + if device_index is None: + device_index = torch.cuda.current_device() if platform.is_cuda_alike() else 0 + caps = platform.get_device_caps(int(device_index)) + provider = OpResolver(DECODE_ATTENTION_LAUNCH_REGISTRY).resolve(spec, caps).provider + return PreparedDecodeAttentionLaunchOp(spec, provider) diff --git a/src/sparsevllm/operators/moe.py b/src/sparsevllm/operators/moe.py index 173e3598..369abe9d 100644 --- a/src/sparsevllm/operators/moe.py +++ b/src/sparsevllm/operators/moe.py @@ -369,7 +369,7 @@ def run( @MOE_REGISTRY.register class FlashInferCutlassFp8MoeProvider(MoeProvider): name = "flashinfer_cutlass_fp8_sm90" - priority = 120 + priority = 100 gate_up_order = "up_gate" @classmethod @@ -433,7 +433,7 @@ def run( output=output, use_deepseek_fp8_block_scale=True, use_fused_finalize=False, - enable_pdl=False, + enable_pdl=None, activation_type=ActivationType.Swiglu, ) return output diff --git a/src/sparsevllm/operators/prefill_attention.py b/src/sparsevllm/operators/prefill_attention.py index f2672b06..921dffb2 100644 --- a/src/sparsevllm/operators/prefill_attention.py +++ b/src/sparsevllm/operators/prefill_attention.py @@ -1,5 +1,6 @@ from __future__ import annotations +import math import re from dataclasses import dataclass from importlib.metadata import PackageNotFoundError, version @@ -45,8 +46,16 @@ class PrefillAttentionProvider: name = "" priority = 0 - def prepare(self, spec: PrefillAttentionOpSpec) -> None: - del spec + def prepare( + self, + spec: PrefillAttentionOpSpec, + *, + device_index: int | None = None, + ) -> None: + del spec, device_index + + def close(self) -> None: + pass def run( self, @@ -62,6 +71,19 @@ def run( raise NotImplementedError +def _validate_token_page_table(view: Any) -> None: + if view.active_slots.dtype != torch.int32: + raise TypeError( + "Paged prefill requires an int32 physical-slot page table, got " + f"{view.active_slots.dtype}." + ) + if view.active_slots.ndim != 2: + raise ValueError( + "Paged prefill expects a 2D physical-slot page table, got " + f"shape={tuple(view.active_slots.shape)}." + ) + + PREFILL_ATTENTION_REGISTRY: OpRegistry[ PrefillAttentionOpSpec, PrefillAttentionProvider ] = OpRegistry("paged prefill attention") @@ -136,20 +158,27 @@ def supports( def __init__(self) -> None: self._state: _FlashInferPagedPrefillState | None = None - def prepare(self, spec: PrefillAttentionOpSpec) -> None: - device = torch.device("cuda", torch.cuda.current_device()) - key = ( - device.index, - spec.num_query_heads, - spec.num_kv_heads, - spec.head_dim, - spec.activation_dtype, - ) - state = _FLASHINFER_PREFILL_STATES.get(key) - if state is None: - state = _FlashInferPagedPrefillState(device) - _FLASHINFER_PREFILL_STATES[key] = state - self._state = state + def prepare( + self, + spec: PrefillAttentionOpSpec, + *, + device_index: int | None = None, + ) -> None: + if self._state is not None: + return + current_device = torch.cuda.current_device() + if device_index is None: + device_index = current_device + if int(device_index) != current_device: + raise RuntimeError( + "FlashInfer prefill must be prepared on the selected CUDA device: " + f"selected={device_index} current={current_device}." + ) + device = torch.device("cuda", int(device_index)) + self._state = _FlashInferPagedPrefillState(device) + + def close(self) -> None: + self._state = None def run( self, @@ -164,8 +193,9 @@ def run( ): del chunk_lens payload, meta = _view_parts(view) + _validate_token_page_table(meta) if self._state is None: - self.prepare(spec) + self.prepare(spec, device_index=q.device.index) state = self._state assert state is not None if q.dtype != spec.activation_dtype: @@ -287,9 +317,6 @@ def plan( self.planned = True -_FLASHINFER_PREFILL_STATES: dict[tuple, _FlashInferPagedPrefillState] = {} - - @PREFILL_ATTENTION_REGISTRY.register class TritonPagedPrefillAttentionProvider(PrefillAttentionProvider): name = "triton_paged_prefill" @@ -309,6 +336,23 @@ def supports( ) if spec.head_dim not in {16, 32, 64, 128, 256}: return SupportResult.no(f"unsupported head_dim={spec.head_dim}") + if not spec.causal: + return SupportResult.no("legacy Triton prefill requires causal attention") + if spec.page_size != 1: + return SupportResult.no( + f"legacy Triton prefill requires page_size=1, got {spec.page_size}" + ) + expected_scale = spec.head_dim**-0.5 + if not math.isclose( + spec.softmax_scale, + expected_scale, + rel_tol=1.0e-6, + abs_tol=0.0, + ): + return SupportResult.no( + "legacy Triton prefill requires the default head-dimension scale " + f"{expected_scale}, got {spec.softmax_scale}" + ) return SupportResult.yes() def run( @@ -324,6 +368,7 @@ def run( ): del spec, layer_idx payload, meta = _view_parts(view) + _validate_token_page_table(meta) from sparsevllm.triton_kernel.context_flashattention_nopad import ( context_attention_fwd, ) @@ -355,3 +400,41 @@ def resolve_prefill_attention_provider( device_index = torch.cuda.current_device() if platform.is_cuda_alike() else 0 caps = platform.get_device_caps(int(device_index)) return OpResolver(PREFILL_ATTENTION_REGISTRY).resolve(spec, caps).provider + + +class PreparedPrefillAttentionOp: + """One prepared provider shared by all layers in one model runtime.""" + + def __init__( + self, + spec: PrefillAttentionOpSpec, + provider: PrefillAttentionProvider, + ) -> None: + self.spec = spec + self.provider = provider + self._closed = False + + @property + def name(self) -> str: + return self.provider.name + + def run(self, q, view, **kwargs): + if self._closed: + raise RuntimeError("Prefill attention operator is closed.") + return self.provider.run(self.spec, q, view, **kwargs) + + def close(self) -> None: + if self._closed: + return + self.provider.close() + self._closed = True + + +def prepare_prefill_attention_op( + spec: PrefillAttentionOpSpec, + *, + device_index: int | None = None, +) -> PreparedPrefillAttentionOp: + provider = resolve_prefill_attention_provider(spec, device_index=device_index) + provider.prepare(spec, device_index=device_index) + return PreparedPrefillAttentionOp(spec, provider) diff --git a/src/sparsevllm/triton_kernel/column_parallel_rmsnorm.py b/src/sparsevllm/triton_kernel/column_parallel_rmsnorm.py new file mode 100644 index 00000000..9cfec1a8 --- /dev/null +++ b/src/sparsevllm/triton_kernel/column_parallel_rmsnorm.py @@ -0,0 +1,230 @@ +"""Fused CUDA kernels for paired column-parallel RMSNorm.""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16) + + +@triton.jit +def _paired_square_sum_kernel( + x_ptr, + other_ptr, + square_sums_ptr, + stride_x_row, + stride_other_row, + stride_sums_row, + x_hidden_size: tl.constexpr, + other_hidden_size: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + row = tl.program_id(0) + cols = tl.arange(0, BLOCK_SIZE) + x_mask = cols < x_hidden_size + other_mask = cols < other_hidden_size + x = tl.load( + x_ptr + row * stride_x_row + cols, + mask=x_mask, + other=0.0, + ).to(tl.float32) + other = tl.load( + other_ptr + row * stride_other_row + cols, + mask=other_mask, + other=0.0, + ).to(tl.float32) + x_sum = tl.sum(tl.where(x_mask, x * x, 0.0), axis=0) + other_sum = tl.sum(tl.where(other_mask, other * other, 0.0), axis=0) + sums_row = square_sums_ptr + row * stride_sums_row + tl.store(sums_row, x_sum) + tl.store(sums_row + 1, other_sum) + + +@triton.jit +def _paired_rms_apply_kernel( + x_ptr, + other_ptr, + square_sums_ptr, + x_weight_ptr, + other_weight_ptr, + x_output_ptr, + other_output_ptr, + stride_x_row, + stride_other_row, + stride_sums_row, + stride_x_output_row, + stride_other_output_row, + x_global_hidden_size: tl.constexpr, + other_global_hidden_size: tl.constexpr, + x_hidden_size: tl.constexpr, + other_hidden_size: tl.constexpr, + x_eps: tl.constexpr, + other_eps: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + row = tl.program_id(0) + cols = tl.arange(0, BLOCK_SIZE) + x_mask = cols < x_hidden_size + other_mask = cols < other_hidden_size + sums_row = square_sums_ptr + row * stride_sums_row + x_inv_rms = tl.rsqrt(tl.load(sums_row) / x_global_hidden_size + x_eps) + other_inv_rms = tl.rsqrt( + tl.load(sums_row + 1) / other_global_hidden_size + other_eps + ) + + x = tl.load( + x_ptr + row * stride_x_row + cols, + mask=x_mask, + other=0.0, + ).to(tl.float32) + other = tl.load( + other_ptr + row * stride_other_row + cols, + mask=other_mask, + other=0.0, + ).to(tl.float32) + x_weight = tl.load(x_weight_ptr + cols, mask=x_mask, other=0.0) + other_weight = tl.load( + other_weight_ptr + cols, + mask=other_mask, + other=0.0, + ) + x_dtype = x_output_ptr.dtype.element_ty + other_dtype = other_output_ptr.dtype.element_ty + x_normalized = (x * x_inv_rms).to(x_dtype) + other_normalized = (other * other_inv_rms).to(other_dtype) + tl.store( + x_output_ptr + row * stride_x_output_row + cols, + x_normalized * x_weight, + mask=x_mask, + ) + tl.store( + other_output_ptr + row * stride_other_output_row + cols, + other_normalized * other_weight, + mask=other_mask, + ) + + +def _validate_pair( + x: torch.Tensor, + other: torch.Tensor, + x_weight: torch.Tensor | None = None, + other_weight: torch.Tensor | None = None, +) -> None: + if not x.is_cuda or not other.is_cuda: + raise ValueError("Paired column-parallel RMSNorm requires CUDA tensors.") + if x.device != other.device or x.dtype != other.dtype: + raise ValueError("Paired RMSNorm inputs must share a device and dtype.") + if x.dtype not in _SUPPORTED_DTYPES: + raise TypeError("Paired RMSNorm supports only FP16 and BF16 inputs.") + if x.ndim != 2 or other.ndim != 2 or x.shape[0] != other.shape[0]: + raise ValueError("Paired RMSNorm inputs must have matching two-dimensional rows.") + if x.shape[1] <= 0 or other.shape[1] <= 0: + raise ValueError("Paired RMSNorm requires non-empty feature dimensions.") + if x.stride(1) != 1 or other.stride(1) != 1: + raise ValueError("Paired RMSNorm requires contiguous feature dimensions.") + if x_weight is None or other_weight is None: + return + for name, weight, hidden_size in ( + ("x", x_weight, x.shape[1]), + ("other", other_weight, other.shape[1]), + ): + if weight.shape != (hidden_size,): + raise ValueError( + f"Paired RMSNorm {name} weight must have shape ({hidden_size},)." + ) + if weight.device != x.device or weight.dtype != x.dtype: + raise ValueError( + f"Paired RMSNorm {name} weight must match the inputs." + ) + if not weight.is_contiguous(): + raise ValueError(f"Paired RMSNorm {name} weight must be contiguous.") + + +def _launch_config(x: torch.Tensor, other: torch.Tensor) -> tuple[int, int]: + hidden_size = max(int(x.shape[1]), int(other.shape[1])) + block_size = triton.next_power_of_2(hidden_size) + if block_size * x.element_size() > 65536: + raise RuntimeError( + "Paired RMSNorm does not support feature dimensions occupying " + "more than 64 KiB per row." + ) + num_warps = min(max(block_size // 256, 1), 8) + return block_size, num_warps + + +def paired_square_sums(x: torch.Tensor, other: torch.Tensor) -> torch.Tensor: + """Return one FP32 square-sum pair per input row.""" + _validate_pair(x, other) + block_size, num_warps = _launch_config(x, other) + square_sums = torch.empty((x.shape[0], 2), dtype=torch.float32, device=x.device) + with torch.cuda.device(x.device): + _paired_square_sum_kernel[(x.shape[0],)]( + x, + other, + square_sums, + x.stride(0), + other.stride(0), + square_sums.stride(0), + x_hidden_size=int(x.shape[1]), + other_hidden_size=int(other.shape[1]), + BLOCK_SIZE=block_size, + num_warps=num_warps, + ) + return square_sums + + +def paired_rms_apply( + x: torch.Tensor, + other: torch.Tensor, + square_sums: torch.Tensor, + x_weight: torch.Tensor, + other_weight: torch.Tensor, + *, + x_global_hidden_size: int, + other_global_hidden_size: int, + x_eps: float, + other_eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + """Apply paired RMSNorm after the statistics have been reduced across TP.""" + _validate_pair(x, other, x_weight, other_weight) + expected_sums_shape = (x.shape[0], 2) + if ( + square_sums.shape != expected_sums_shape + or square_sums.dtype != torch.float32 + or square_sums.device != x.device + or not square_sums.is_contiguous() + ): + raise ValueError( + "Paired RMSNorm square sums must be contiguous FP32 values with " + f"shape {expected_sums_shape} on the input device." + ) + block_size, num_warps = _launch_config(x, other) + x_output = torch.empty_like(x, memory_format=torch.contiguous_format) + other_output = torch.empty_like(other, memory_format=torch.contiguous_format) + with torch.cuda.device(x.device): + _paired_rms_apply_kernel[(x.shape[0],)]( + x, + other, + square_sums, + x_weight, + other_weight, + x_output, + other_output, + x.stride(0), + other.stride(0), + square_sums.stride(0), + x_output.stride(0), + other_output.stride(0), + x_global_hidden_size=int(x_global_hidden_size), + other_global_hidden_size=int(other_global_hidden_size), + x_hidden_size=int(x.shape[1]), + other_hidden_size=int(other.shape[1]), + x_eps=float(x_eps), + other_eps=float(other_eps), + BLOCK_SIZE=block_size, + num_warps=num_warps, + ) + return x_output, other_output diff --git a/tests/test_column_parallel_rmsnorm.py b/tests/test_column_parallel_rmsnorm.py index 0451cd38..159b2d82 100644 --- a/tests/test_column_parallel_rmsnorm.py +++ b/tests/test_column_parallel_rmsnorm.py @@ -1,5 +1,6 @@ from types import SimpleNamespace +import pytest import torch from sparsevllm.layers.layernorm import ColumnParallelRMSNorm @@ -64,3 +65,50 @@ def test_column_parallel_rmsnorm_exposes_rank_local_weight_slice(): norm = ColumnParallelRMSNorm(8, parallel_context=context) assert norm.rank_local_weight_slice((8,)) == (slice(4, 8),) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_column_parallel_qk_norm_cuda_fused_path_matches_reference(): + torch.manual_seed(17) + query = torch.randn(5, 16, device="cuda", dtype=torch.bfloat16) + key = torch.randn(5, 8, device="cuda", dtype=torch.bfloat16) + query_weight = torch.randn(16, device="cuda", dtype=torch.bfloat16) + key_weight = torch.randn(8, device="cuda", dtype=torch.bfloat16) + outputs = [] + + for rank in range(2): + local_query = query.chunk(2, dim=-1)[rank] + local_key = key.chunk(2, dim=-1)[rank] + other_rank = 1 - rank + remote_sums = torch.stack( + ( + query.chunk(2, dim=-1)[other_rank].float().square().sum(-1), + key.chunk(2, dim=-1)[other_rank].float().square().sum(-1), + ), + dim=-1, + ) + context = _ReferenceTpContext(rank, remote_sums) + q_norm = ColumnParallelRMSNorm(16, parallel_context=context).to( + device="cuda", dtype=torch.bfloat16 + ) + k_norm = ColumnParallelRMSNorm(8, parallel_context=context).to( + device="cuda", dtype=torch.bfloat16 + ) + q_norm.weight.data.copy_(query_weight.chunk(2)[rank]) + k_norm.weight.data.copy_(key_weight.chunk(2)[rank]) + + outputs.append(q_norm.forward_pair(local_query, local_key, k_norm)) + assert context.all_reduce_calls == 1 + + actual_query = torch.cat([item[0] for item in outputs], dim=-1) + actual_key = torch.cat([item[1] for item in outputs], dim=-1) + expected_query = ( + query.float() + * torch.rsqrt(query.float().square().mean(-1, keepdim=True) + 1.0e-6) + ).to(query.dtype) * query_weight + expected_key = ( + key.float() * torch.rsqrt(key.float().square().mean(-1, keepdim=True) + 1.0e-6) + ).to(key.dtype) * key_weight + + torch.testing.assert_close(actual_query, expected_query, atol=0.015625, rtol=0) + torch.testing.assert_close(actual_key, expected_key, atol=0.015625, rtol=0) diff --git a/tests/test_minimax_m2_model.py b/tests/test_minimax_m2_model.py index 28fb8599..620ae39f 100644 --- a/tests/test_minimax_m2_model.py +++ b/tests/test_minimax_m2_model.py @@ -1,5 +1,5 @@ from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import Mock, patch import pytest import torch @@ -12,6 +12,7 @@ MiniMaxM2Attention, MiniMaxM2ForCausalLM, MiniMaxM2PackedExperts, + MiniMaxM2RuntimeConfig, ) from sparsevllm.utils.loader import load_model @@ -147,7 +148,7 @@ def _tp_context(tp_rank: int, tp_size: int) -> ParallelContext: ) -def _instantiate_model(config, context): +def _instantiate_model(config, context, runtime_config=None): with ( patch( "sparsevllm.models.minimax_m2.get_parallel_context", @@ -166,7 +167,38 @@ def _instantiate_model(config, context): return_value=context, ), ): - return MiniMaxM2ForCausalLM(config) + return MiniMaxM2ForCausalLM(config, runtime_config=runtime_config) + + +def test_model_construction_shares_explicit_runtime_operators_across_layers(): + config = _config(num_hidden_layers=2) + prefill_op = SimpleNamespace(name="prefill", close=Mock()) + decode_op = SimpleNamespace(name="decode") + all_reduce_op = SimpleNamespace(name="all_reduce", run=Mock(), close=Mock()) + runtime_config = MiniMaxM2RuntimeConfig( + prefill_attention_op=prefill_op, + decode_launch_op=decode_op, + attention_decode_all_reduce=all_reduce_op, + moe_decode_all_reduce=all_reduce_op, + cuda_graph=True, + ) + + model = _instantiate_model( + config, + _tp_context(0, 1), + runtime_config=runtime_config, + ) + + assert not hasattr(config, "prefill_kv_view_layer_invariant") + for layer in model.model.layers: + assert layer.self_attn.attn.prefill_op is prefill_op + assert layer.self_attn.attn.decode_launch_op is decode_op + assert not layer.self_attn.o_proj.reduce_results + assert layer.block_sparse_moe.experts.op_spec.cuda_graph + model.close_runtime_operators() + model.close_runtime_operators() + prefill_op.close.assert_called_once_with() + all_reduce_op.close.assert_called_once_with() def _random_fp8(shape): diff --git a/tests/test_parallel_context.py b/tests/test_parallel_context.py index 30865b83..1ee6e1f8 100644 --- a/tests/test_parallel_context.py +++ b/tests/test_parallel_context.py @@ -393,6 +393,28 @@ def test_all_reduce_capture_and_close_delegate_to_provider(): provider.close.assert_called_once_with() +@pytest.mark.parametrize("op", [dist.ReduceOp.SUM, dist.ReduceOp.MAX]) +def test_parallel_context_collectives_are_always_in_place_torch_operations(op): + world_group = object() + context = ParallelContext( + world=ParallelGroup(world_group, (0, 1, 2, 3), 0, 4), + tensor=ParallelGroup(world_group, (0, 1, 2, 3), 0, 4), + expert=ParallelGroup(None, (0,), 0, 1), + data=ParallelGroup(None, (0,), 0, 1), + ) + tensor = torch.ones(2, 3072, dtype=torch.bfloat16) + + with patch.object(dist, "all_reduce") as all_reduce: + returned = context.world_all_reduce(tensor, op=op) + + assert returned is tensor + all_reduce.assert_called_once_with( + tensor, + op=op, + group=world_group, + ) + + def test_qwen3_moe_parallel_config_validation(tmp_path): with patch("sparsevllm.configs.runtime.AutoConfig.from_pretrained", return_value=_hf_config()): config = Config(model=str(tmp_path), expert_parallel_size=4) diff --git a/tests/test_prefill_attention_provider.py b/tests/test_prefill_attention_provider.py index cf8b63ee..6671893e 100644 --- a/tests/test_prefill_attention_provider.py +++ b/tests/test_prefill_attention_provider.py @@ -1,5 +1,7 @@ +import gc +import weakref from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import Mock, patch import pytest import torch @@ -7,7 +9,9 @@ from sparsevllm.operators.prefill_attention import ( PREFILL_ATTENTION_REGISTRY, FlashInferPagedPrefillAttentionProvider, + PreparedPrefillAttentionOp, PrefillAttentionOpSpec, + TritonPagedPrefillAttentionProvider, ) from sparsevllm.operators.registry import OpResolver from sparsevllm.platforms import DeviceCaps, PlatformEnum @@ -103,6 +107,103 @@ def test_resolver_falls_back_when_flashinfer_is_too_old(_find, _version): ) in resolved.rejected +@pytest.mark.parametrize( + ("spec", "reason"), + [ + (_spec(causal=False), "causal attention"), + (_spec(page_size=16), "page_size=1"), + (_spec(softmax_scale=0.125), "default head-dimension scale"), + ], +) +def test_triton_provider_rejects_unimplemented_attention_semantics(spec, reason): + result = TritonPagedPrefillAttentionProvider.supports(spec, _h100_caps()) + + assert not result.supported + assert reason in result.reason + + +@patch( + "sparsevllm.operators.prefill_attention.version", + return_value="0.6.15", +) +@patch("sparsevllm.operators.prefill_attention.find_spec", return_value=object()) +def test_resolver_rejects_page_sizes_without_a_valid_kernel(_find, _version): + with pytest.raises(RuntimeError, match="No paged prefill attention provider"): + OpResolver(PREFILL_ATTENTION_REGISTRY).resolve( + _spec(page_size=16), + _h100_caps(), + ) + + +def test_paged_prefill_rejects_non_int32_page_table_before_kernel_launch(): + provider = TritonPagedPrefillAttentionProvider() + view = SimpleNamespace(active_slots=torch.zeros(1, 2, dtype=torch.int64)) + + with pytest.raises(TypeError, match="int32 physical-slot page table"): + provider.run( + _spec(), + torch.empty(1, 12, 128), + view, + qo_indptr=torch.tensor([0, 1], dtype=torch.int32), + chunk_lens=torch.tensor([1], dtype=torch.int32), + max_context_len=1, + layer_idx=0, + ) + + +def test_prepared_prefill_close_releases_provider_state(): + class State: + pass + + provider = FlashInferPagedPrefillAttentionProvider() + state = State() + state_ref = weakref.ref(state) + provider._state = state + op = PreparedPrefillAttentionOp(_spec(), provider) + del state + + op.close() + gc.collect() + + assert provider._state is None + assert state_ref() is None + with pytest.raises(RuntimeError, match="closed"): + op.run(None, None) + + +def test_flashinfer_prefill_passes_kv_cache_as_page_views_without_copying(): + wrapper = SimpleNamespace(run=Mock()) + provider = FlashInferPagedPrefillAttentionProvider() + provider._state = SimpleNamespace(planned=True, wrapper=wrapper) + q = torch.empty(1, 12, 128, dtype=torch.bfloat16) + k_cache = torch.empty(5, 2, 128, dtype=torch.bfloat16) + v_cache = torch.empty_like(k_cache) + view = SimpleNamespace( + k_cache=k_cache, + v_cache=v_cache, + active_slots=torch.tensor([[4, 1]], dtype=torch.int32), + req_indices=torch.tensor([0], dtype=torch.int32), + context_lens=torch.tensor([2], dtype=torch.int32), + attn_score=None, + ) + + provider.run( + _spec(), + q, + view, + qo_indptr=torch.tensor([0, 1], dtype=torch.int32), + chunk_lens=torch.tensor([1], dtype=torch.int32), + max_context_len=2, + layer_idx=1, + ) + + paged_k, paged_v = wrapper.run.call_args.args[1] + assert paged_k.shape == (5, 1, 2, 128) + assert paged_v.shape == (5, 1, 2, 128) + assert paged_k.data_ptr() == k_cache.data_ptr() + assert paged_v.data_ptr() == v_cache.data_ptr() + + def _torch_prefill_oracle(q, logical_k, logical_v, q_lens, kv_lens): outputs = [] q_cursor = 0 @@ -162,3 +263,47 @@ def test_flashinfer_page_size_one_matches_noncontiguous_torch_oracle(): q, logical_k, logical_v, q_lens, kv_lens ) torch.testing.assert_close(actual, expected, rtol=0.03, atol=0.03) + provider.close() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_triton_page_size_one_matches_noncontiguous_torch_oracle(): + torch.manual_seed(20260809) + q_lens = [3, 2] + kv_lens = [5, 6] + q = torch.randn(5, 12, 128, device="cuda", dtype=torch.bfloat16) + k_cache = torch.randn(23, 2, 128, device="cuda", dtype=torch.bfloat16) + v_cache = torch.randn_like(k_cache) + pages = torch.randperm(23, device="cuda")[:11] + rows = torch.zeros(2, 6, device="cuda", dtype=torch.int32) + rows[0, :5] = pages[:5].to(torch.int32) + rows[1, :6] = pages[5:].to(torch.int32) + logical_k = [k_cache[pages[:5]], k_cache[pages[5:]]] + logical_v = [v_cache[pages[:5]], v_cache[pages[5:]]] + view = SimpleNamespace( + k_cache=k_cache, + v_cache=v_cache, + active_slots=rows, + req_indices=torch.tensor([0, 1], device="cuda", dtype=torch.int32), + context_lens=torch.tensor(kv_lens, device="cuda", dtype=torch.int32), + attn_score=None, + ) + provider = TritonPagedPrefillAttentionProvider() + + actual = provider.run( + _spec(), + q, + view, + qo_indptr=torch.tensor([0, 3, 5], device="cuda", dtype=torch.int32), + chunk_lens=torch.tensor(q_lens, device="cuda", dtype=torch.int32), + max_context_len=6, + layer_idx=0, + ) + expected = _torch_prefill_oracle( + q, + logical_k, + logical_v, + q_lens, + kv_lens, + ) + torch.testing.assert_close(actual, expected, rtol=0.03, atol=0.03) diff --git a/tests/test_prefill_schedule_policy.py b/tests/test_prefill_schedule_policy.py index 025f024f..16652274 100644 --- a/tests/test_prefill_schedule_policy.py +++ b/tests/test_prefill_schedule_policy.py @@ -11,6 +11,7 @@ import torch from sparsevllm.config import Config +from sparsevllm.configs.cuda_graph import _resolve_decode_static_batch_capacity from sparsevllm.engine.cache_manager.standard import StandardCacheManager from sparsevllm.engine.cache_manager.deltakv import DeltaKVCacheManager from sparsevllm.engine.cache_manager.deltakv_less_memory import DeltaKVLessMemoryCacheManager @@ -1140,6 +1141,27 @@ def test_decode_cuda_graph_auto_capture_sizes_end_at_decode_limit(self): self.assertTrue(cfg.decode_graph) self.assertEqual(cfg.decode_graph_capture_sizes, expected_sizes) + def test_decode_static_batch_capacity_uses_reachable_padding_bucket(self): + cases = ( + ([1, 2, 4, 8, 16, 32, 64], 32, 64, 32), + ([1, 4, 8, 64], 32, 64, 64), + ([1, 2, 4, 8, 16, 32, 64], 80, 64, 64), + ) + for capture_sizes, max_batch, max_decode, expected in cases: + with self.subTest( + capture_sizes=capture_sizes, + max_batch=max_batch, + max_decode=max_decode, + ): + self.assertEqual( + _resolve_decode_static_batch_capacity( + capture_sizes, + max_num_seqs_in_batch=max_batch, + max_decoding_seqs=max_decode, + ), + expected, + ) + def test_decode_graph_aliases_normalize_to_canonical_fields(self): cfg = self.make_config( vllm_sparse_method="omnikv", diff --git a/tests/test_qwen35_hd256_decode_routing.py b/tests/test_qwen35_hd256_decode_routing.py index 642cbd06..09daeb27 100644 --- a/tests/test_qwen35_hd256_decode_routing.py +++ b/tests/test_qwen35_hd256_decode_routing.py @@ -17,43 +17,6 @@ class Qwen35Hd256DecodeRoutingTest(unittest.TestCase): - def test_minimax_long_decode_uses_tuned_launch_config(self): - config = TritonAttentionBackend.gqa_decode_launch_config( - block_seq=256, - max_context_len=65536, - num_heads=12, - num_kv_heads=2, - head_dim=128, - requires_attention_scores=False, - ) - - self.assertEqual(config, (1024, 128, 4)) - - def test_other_gqa_decode_shapes_keep_default_launch_config(self): - cases = ( - {"max_context_len": 32768}, - {"num_heads": 16}, - {"num_kv_heads": 4}, - {"head_dim": 256}, - {"requires_attention_scores": True}, - {"block_seq": 128}, - ) - defaults = { - "block_seq": 256, - "max_context_len": 65536, - "num_heads": 12, - "num_kv_heads": 2, - "head_dim": 128, - "requires_attention_scores": False, - } - - for override in cases: - with self.subTest(override=override): - config = TritonAttentionBackend.gqa_decode_launch_config( - **(defaults | override) - ) - self.assertEqual(config, (override.get("block_seq", 256), 16, 2)) - def _make_view(self, *, head_dim: int, attn_score=None): active_slots = torch.tensor([[0, 1, 2]], dtype=torch.int32) req_indices = torch.tensor([0], dtype=torch.int32) diff --git a/tests/test_tp_rpc.py b/tests/test_tp_rpc.py index 31343c53..65cb8f83 100644 --- a/tests/test_tp_rpc.py +++ b/tests/test_tp_rpc.py @@ -443,6 +443,9 @@ def test_model_runner_exit_drains_graphs_before_barrier(): runner.decode_cuda_graph_runner = SimpleNamespace( clear_captured_graphs=lambda: calls.append("clear_graphs") ) + runner.model = SimpleNamespace( + close_runtime_operators=lambda: calls.append("close_ops") + ) runner.world_size = 2 runner.rank = 0 runner.shm = SimpleNamespace( @@ -470,6 +473,8 @@ def test_model_runner_exit_drains_graphs_before_barrier(): "sync", "clear_graphs", "sync", + "close_ops", + "sync", "close_shm", "barrier", "unlink_shm", From 73f418910afb52475adfbfa2656568ca04c9f8b1 Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Wed, 12 Aug 2026 16:13:10 +0800 Subject: [PATCH 25/48] refactor: align glm provider contracts --- src/sparsevllm/configs/model.py | 65 +----- src/sparsevllm/engine/cache_manager/base.py | 2 +- src/sparsevllm/engine/sparse_controller.py | 2 +- src/sparsevllm/models/glm4_moe_lite.py | 2 +- src/sparsevllm/models/layout.py | 23 +++ src/sparsevllm/models/spec.py | 2 + src/sparsevllm/operators/moe.py | 10 +- src/sparsevllm/operators/sgl_fa3.py | 102 +--------- src/sparsevllm/operators/sgl_kernel.py | 25 +++ src/sparsevllm/operators/sgl_moe.py | 209 ++------------------ tests/test_glm_cuda_graph.py | 2 +- tests/test_glm_runtime_compatibility.py | 16 -- tests/test_operator_providers.py | 2 +- tests/test_qwen35_mixed_runtime.py | 4 + tests/test_sgl_fa3.py | 27 +-- tests/test_sgl_moe.py | 107 ++-------- 16 files changed, 105 insertions(+), 495 deletions(-) create mode 100644 src/sparsevllm/operators/sgl_kernel.py diff --git a/src/sparsevllm/configs/model.py b/src/sparsevllm/configs/model.py index 2cb3503e..3e68ca29 100644 --- a/src/sparsevllm/configs/model.py +++ b/src/sparsevllm/configs/model.py @@ -5,7 +5,6 @@ from types import SimpleNamespace from typing import Any -import torch from transformers import AutoConfig from sparsevllm.method_registry import ( @@ -24,30 +23,6 @@ def _config_to_namespace(config: dict[str, Any]) -> SimpleNamespace: return SimpleNamespace(**config) -def resolve_attention_qk_head_dim(hf_config: Any) -> int: - qk_nope_head_dim = config_get(hf_config, "qk_nope_head_dim", None) - qk_rope_head_dim = config_get(hf_config, "qk_rope_head_dim", None) - if qk_nope_head_dim is not None and qk_rope_head_dim is not None: - head_dim = int(qk_nope_head_dim) + int(qk_rope_head_dim) - if head_dim <= 0: - raise ValueError(f"Attention QK head dimension must be positive, got {head_dim}.") - return head_dim - head_dim = config_get(hf_config, "head_dim", None) - if head_dim is not None: - head_dim = int(head_dim) - if head_dim <= 0: - raise ValueError(f"Attention QK head dimension must be positive, got {head_dim}.") - return head_dim - hidden_size = int(config_get(hf_config, "hidden_size", 0) or 0) - num_heads = int(config_get(hf_config, "num_attention_heads", 0) or 0) - if hidden_size <= 0 or num_heads <= 0 or hidden_size % num_heads: - raise ValueError( - "Attention QK head dimension requires valid head_dim or divisible " - "hidden_size/num_attention_heads." - ) - return hidden_size // num_heads - - def _load_model_config(model_path: str) -> Any: try: return AutoConfig.from_pretrained(model_path, trust_remote_code=True) @@ -87,44 +62,6 @@ def _extract_text_config(config: Any) -> Any: return text_config -def _resolve_attention_cache_layout(config, model_type: str) -> None: - config.attention_cache_layout = ( - "mla_latent" if model_type == "glm4_moe_lite" else "explicit_kv" - ) - if config.attention_cache_layout == "explicit_kv": - return - if config.enable_prefix_cache_offload: - raise NotImplementedError( - "GLM-4.7-Flash latent MLA does not support prefix cache offload." - ) - if config.quantization_config.enabled: - raise NotImplementedError( - "GLM-4.7-Flash latent MLA does not support quantized checkpoints." - ) - if config.hf_config.torch_dtype != torch.bfloat16: - raise NotImplementedError( - "GLM-4.7-Flash latent MLA requires BF16 checkpoints." - ) - expected = { - "num_attention_heads": 20, - "q_lora_rank": 768, - "kv_lora_rank": 512, - "qk_nope_head_dim": 192, - "qk_rope_head_dim": 64, - "v_head_dim": 256, - } - invalid = { - field: config_get(config.hf_config, field, None) - for field, value in expected.items() - if config_get(config.hf_config, field, None) != value - } - if invalid: - field, actual = next(iter(invalid.items())) - raise NotImplementedError( - f"GLM-4.7-Flash latent MLA requires {field}={expected[field]}, got {actual}." - ) - - def _finalize_model_config(config, model_spec: ModelSpec) -> None: config.runtime_layout = RuntimeLayout.from_config( config.hf_config, @@ -195,7 +132,7 @@ def load_and_validate_model(config) -> None: "Tiny random mode does not support quantized model weights." ) setattr(config.hf_config, "quantization_config", config.quantization_config) - _resolve_attention_cache_layout(config, model_type) + config.attention_cache_layout = model_spec.attention_cache_layout validate_checkpoint( model_type, outer_config=config.outer_hf_config, diff --git a/src/sparsevllm/engine/cache_manager/base.py b/src/sparsevllm/engine/cache_manager/base.py index 7dba1f86..0b85bf0d 100644 --- a/src/sparsevllm/engine/cache_manager/base.py +++ b/src/sparsevllm/engine/cache_manager/base.py @@ -20,7 +20,7 @@ from sparsevllm.method_registry import SUPPORTED_SPARSE_METHODS, normalize_sparse_method from sparsevllm.triton_kernel.store_kvcache import store_kvcache import sparsevllm.platforms as platforms -from sparsevllm.configs.model import resolve_attention_qk_head_dim +from sparsevllm.models.layout import resolve_attention_qk_head_dim from sparsevllm.utils.log import logger, log_level diff --git a/src/sparsevllm/engine/sparse_controller.py b/src/sparsevllm/engine/sparse_controller.py index 0a349bff..4406aa9b 100644 --- a/src/sparsevllm/engine/sparse_controller.py +++ b/src/sparsevllm/engine/sparse_controller.py @@ -3,7 +3,7 @@ import torch import torch.nn.functional as F from sparsevllm.config import Config -from sparsevllm.configs.model import resolve_attention_qk_head_dim +from sparsevllm.models.layout import resolve_attention_qk_head_dim from sparsevllm.engine.activation_controller import ActivationController from sparsevllm.engine.sequence import Sequence from sparsevllm.engine.cache_manager import CacheManager, SparseSelection diff --git a/src/sparsevllm/models/glm4_moe_lite.py b/src/sparsevllm/models/glm4_moe_lite.py index c71adca8..b8516598 100644 --- a/src/sparsevllm/models/glm4_moe_lite.py +++ b/src/sparsevllm/models/glm4_moe_lite.py @@ -9,7 +9,7 @@ from transformers import Glm4MoeLiteConfig from sparsevllm.distributed import get_parallel_context -from sparsevllm.configs.model import resolve_attention_qk_head_dim +from sparsevllm.models.layout import resolve_attention_qk_head_dim from sparsevllm.engine.cache_manager import ( AttentionKeyComputeView, MlaLatentWrite, diff --git a/src/sparsevllm/models/layout.py b/src/sparsevllm/models/layout.py index f1b7eb9d..2e3f450a 100644 --- a/src/sparsevllm/models/layout.py +++ b/src/sparsevllm/models/layout.py @@ -6,6 +6,29 @@ from sparsevllm.utils.config import config_get +def resolve_attention_qk_head_dim(hf_config: Any) -> int: + dimensions = ( + config_get(hf_config, "qk_nope_head_dim", None), + config_get(hf_config, "qk_rope_head_dim", None), + ) + if all(value is not None for value in dimensions): + head_dim = sum(map(int, dimensions)) + elif (value := config_get(hf_config, "head_dim", None)) is not None: + head_dim = int(value) + else: + hidden_size = int(config_get(hf_config, "hidden_size", 0) or 0) + num_heads = int(config_get(hf_config, "num_attention_heads", 0) or 0) + if hidden_size <= 0 or num_heads <= 0 or hidden_size % num_heads: + raise ValueError( + "Attention QK head dimension requires valid head_dim or divisible " + "hidden_size/num_attention_heads." + ) + head_dim = hidden_size // num_heads + if head_dim <= 0: + raise ValueError(f"Attention QK head dimension must be positive, got {head_dim}.") + return head_dim + + def _coerce_int_list( name: str, value: Any, diff --git a/src/sparsevllm/models/spec.py b/src/sparsevllm/models/spec.py index dc4e9892..0e1fae8a 100644 --- a/src/sparsevllm/models/spec.py +++ b/src/sparsevllm/models/spec.py @@ -21,6 +21,7 @@ class ModelSpec: prefix_cache_block_size_multiple: int | None = None deltakv_checkpoint_model_types: frozenset[str] = frozenset() runtime_class_name: str = "" + attention_cache_layout: str = "explicit_kv" attention_tp_fields: tuple[str, ...] = () num_experts_field: str | None = None moe_tp_fields: tuple[str, ...] = () @@ -164,6 +165,7 @@ def validate_sharding(self, hf_config: Any, topology: ParallelTopology) -> None: supports_expert_parallel=True, supports_outer_tp_moe=True, runtime_class_name="Glm4MoeLiteForCausalLM", + attention_cache_layout="mla_latent", attention_tp_fields=("num_attention_heads", "vocab_size"), num_experts_field="n_routed_experts", moe_tp_fields=("intermediate_size", "moe_intermediate_size"), diff --git a/src/sparsevllm/operators/moe.py b/src/sparsevllm/operators/moe.py index 369abe9d..1804874e 100644 --- a/src/sparsevllm/operators/moe.py +++ b/src/sparsevllm/operators/moe.py @@ -211,15 +211,9 @@ def supports(cls, spec: MoeOpSpec, caps: DeviceCaps) -> SupportResult: return SupportResult.no("requires BF16 activations") if spec.weight_dtype != torch.bfloat16 or spec.block_shape is not None: return SupportResult.no("requires unquantized BF16 expert weights") - from sparsevllm.operators.sgl_moe import ( - sgl_moe_alignment_support, - sgl_moe_ep_alignment_support, - ) + from sparsevllm.operators.sgl_moe import sgl_moe_alignment_support - if spec.ep_size > 1: - supported, reason = sgl_moe_ep_alignment_support() - else: - supported, reason = sgl_moe_alignment_support() + supported, reason = sgl_moe_alignment_support() return SupportResult.yes() if supported else SupportResult.no(reason) def run( diff --git a/src/sparsevllm/operators/sgl_fa3.py b/src/sparsevllm/operators/sgl_fa3.py index e2f31807..be378924 100644 --- a/src/sparsevllm/operators/sgl_fa3.py +++ b/src/sparsevllm/operators/sgl_fa3.py @@ -1,19 +1,11 @@ from __future__ import annotations -import importlib.metadata -import importlib.util -import re from collections.abc import Sequence from dataclasses import dataclass import torch - -_SGL_KERNEL_DISTRIBUTIONS = ("sglang-kernel", "sgl-kernel") -_MIN_VALIDATED_VERSIONS = { - (0, 3): (0, 3, 14), - (0, 4): (0, 4, 5), -} +from sparsevllm.operators.sgl_kernel import sgl_kernel_support _COMMON_FWD_ARGUMENTS = ( "q", @@ -53,10 +45,7 @@ "sm_margin", "sinks", ) -_FWD_04_EXTRA_ARGUMENTS = ( - "sparse_mask_fine", - "only_qv", -) +_FWD_ARGUMENTS = _COMMON_FWD_ARGUMENTS + ("attention_chunk",) + _FWD_TRAILING_ARGUMENTS @dataclass(frozen=True, slots=True) @@ -109,57 +98,8 @@ def matches( ) -def _parse_version(version: str) -> tuple[int, int, int] | None: - parts = str(version).split(".") - if len(parts) < 3: - return None - parsed: list[int] = [] - for part in parts[:3]: - match = re.match(r"(\d+)", part) - if match is None: - return None - parsed.append(int(match.group(1))) - return tuple(parsed) # type: ignore[return-value] - - -def _installed_sgl_kernel_version() -> tuple[str, tuple[int, int, int]] | None: - for distribution in _SGL_KERNEL_DISTRIBUTIONS: - try: - version = importlib.metadata.version(distribution) - except importlib.metadata.PackageNotFoundError: - continue - parsed = _parse_version(version) - if parsed is None: - return version, (-1, -1, -1) - return version, parsed - return None - - def sgl_fa3_support() -> tuple[bool, str]: - """Probe package metadata without importing optional SGL Python wrappers.""" - - try: - module_spec = importlib.util.find_spec("sgl_kernel") - except (ImportError, ValueError): - module_spec = None - if module_spec is None: - return False, "sgl-kernel is not installed" - installed = _installed_sgl_kernel_version() - if installed is None: - return False, "sgl-kernel package metadata is unavailable" - version, parsed = installed - minimum = _MIN_VALIDATED_VERSIONS.get(parsed[:2]) - if minimum is None: - return ( - False, - "requires validated sgl-kernel 0.3.x or sglang-kernel 0.4.x " - f"API, got {version}", - ) - if parsed < minimum: - package = "sglang-kernel" if parsed[:2] == (0, 4) else "sgl-kernel" - required = ".".join(str(value) for value in minimum) - return False, f"requires {package} >= {required}, got {version}" - return True, f"SGL kernel {version} FA3 raw op is available" + return sgl_kernel_support("FA3 raw op") class SglFa3DecodeKernel: @@ -184,23 +124,7 @@ def __init__( op = torch.ops.sgl_kernel.fwd.default argument_names = tuple(argument.name for argument in op._schema.arguments) - legacy = _COMMON_FWD_ARGUMENTS + _FWD_TRAILING_ARGUMENTS - chunked = ( - _COMMON_FWD_ARGUMENTS - + ("attention_chunk",) - + _FWD_TRAILING_ARGUMENTS - ) - chunked_04 = chunked + _FWD_04_EXTRA_ARGUMENTS - if argument_names == legacy: - self._has_attention_chunk = False - self._has_04_arguments = False - elif argument_names == chunked: - self._has_attention_chunk = True - self._has_04_arguments = False - elif argument_names == chunked_04: - self._has_attention_chunk = True - self._has_04_arguments = True - else: + if argument_names != _FWD_ARGUMENTS: raise RuntimeError( "Unsupported sgl-kernel FA3 fwd schema: " f"arguments={argument_names}." @@ -211,9 +135,6 @@ def __init__( "get_scheduler_metadata", None, ) - # sgl-kernel 0.3.14 exposes the FA3 scheduler only internally. Keep - # its validated per-call scheduling path; 0.4.5 exposes the reusable - # raw op used below. self._scheduler_op = ( None if scheduler_packet is None else scheduler_packet.default ) @@ -425,13 +346,10 @@ def run_varlen( -1, -1, ] - if self._has_attention_chunk: - args.append(0) + args.append(0) args.extend( (0.0, True, scheduler_metadata, split_count, None, 0, None) ) - if self._has_04_arguments: - args.extend((None, False)) result: Sequence[torch.Tensor] = self._op(*args) if not result or result[0].data_ptr() != output.data_ptr(): raise RuntimeError("sgl-kernel FA3 did not write to the supplied output") @@ -494,13 +412,10 @@ def run_explicit_varlen( -1, -1, ] - if self._has_attention_chunk: - args.append(0) + args.append(0) args.extend( (0.0, True, scheduler_metadata, self.num_splits, None, 0, None) ) - if self._has_04_arguments: - args.extend((None, False)) result: Sequence[torch.Tensor] = self._op(*args) if not result or result[0].data_ptr() != output.data_ptr(): raise RuntimeError("sgl-kernel FA3 did not write to the supplied output") @@ -549,11 +464,8 @@ def run_contiguous_explicit_varlen( -1, -1, ] - if self._has_attention_chunk: - args.append(0) + args.append(0) args.extend((0.0, True, None, self.num_splits, None, 0, None)) - if self._has_04_arguments: - args.extend((None, False)) result: Sequence[torch.Tensor] = self._op(*args) if not result or result[0].data_ptr() != output.data_ptr(): raise RuntimeError("sgl-kernel FA3 did not write to the supplied output") diff --git a/src/sparsevllm/operators/sgl_kernel.py b/src/sparsevllm/operators/sgl_kernel.py new file mode 100644 index 00000000..30ab6a57 --- /dev/null +++ b/src/sparsevllm/operators/sgl_kernel.py @@ -0,0 +1,25 @@ +from __future__ import annotations + +import importlib.metadata +import importlib.util +import re + + +_MIN_VERSION = (0, 3, 21) +_MAX_VERSION = (0, 4, 0) + + +def sgl_kernel_support(feature: str) -> tuple[bool, str]: + """Check the sgl-kernel API range declared by pyproject.toml.""" + + if importlib.util.find_spec("sgl_kernel") is None: + return False, "sgl-kernel is not installed" + try: + version = importlib.metadata.version("sgl-kernel") + except importlib.metadata.PackageNotFoundError: + return False, "sgl-kernel package metadata is unavailable" + match = re.match(r"^(\d+)\.(\d+)\.(\d+)", version) + parsed = tuple(map(int, match.groups())) if match else None + if parsed is None or not _MIN_VERSION <= parsed < _MAX_VERSION: + return False, f"requires sgl-kernel>=0.3.21,<0.4, got {version}" + return True, f"sgl-kernel {version} {feature} is available" diff --git a/src/sparsevllm/operators/sgl_moe.py b/src/sparsevllm/operators/sgl_moe.py index d574cbbb..85b7cee8 100644 --- a/src/sparsevllm/operators/sgl_moe.py +++ b/src/sparsevllm/operators/sgl_moe.py @@ -1,169 +1,16 @@ from __future__ import annotations -import importlib.metadata -import importlib.util -import re - import torch import triton +from sparsevllm.operators.sgl_kernel import sgl_kernel_support from sparsevllm.triton_kernel.moe import MoeAlignment -_SGL_KERNEL_DISTRIBUTIONS = ("sglang-kernel", "sgl-kernel") -_MIN_VALIDATED_VERSIONS = { - (0, 3): (0, 3, 14), - (0, 4): (0, 4, 5), -} - - -def _parse_version(version: str) -> tuple[int, int, int] | None: - parts = str(version).split(".") - if len(parts) < 3: - return None - parsed: list[int] = [] - for part in parts[:3]: - match = re.match(r"(\d+)", part) - if match is None: - return None - parsed.append(int(match.group(1))) - return tuple(parsed) # type: ignore[return-value] - - -def _installed_sgl_kernel_version() -> tuple[str, tuple[int, int, int]] | None: - for distribution in _SGL_KERNEL_DISTRIBUTIONS: - try: - version = importlib.metadata.version(distribution) - except importlib.metadata.PackageNotFoundError: - continue - parsed = _parse_version(version) - if parsed is None: - return version, (-1, -1, -1) - return version, parsed - return None - - -def _sgl_moe_support(feature: str) -> tuple[bool, str]: - try: - module_spec = importlib.util.find_spec("sgl_kernel") - except (ImportError, ValueError): - module_spec = None - if module_spec is None: - return False, "sgl-kernel is not installed" - installed = _installed_sgl_kernel_version() - if installed is None: - return False, "sgl-kernel package metadata is unavailable" - version, parsed = installed - minimum = _MIN_VALIDATED_VERSIONS.get(parsed[:2]) - if minimum is None: - return ( - False, - "requires validated sgl-kernel 0.3.x or sglang-kernel 0.4.x " - f"API, got {version}", - ) - if parsed < minimum: - package = "sglang-kernel" if parsed[:2] == (0, 4) else "sgl-kernel" - required = ".".join(str(value) for value in minimum) - return False, f"requires {package} >= {required}, got {version}" - return True, f"SGL kernel {version} {feature} is available" - - -def sgl_fused_moe_gate_support() -> tuple[bool, str]: - """Check the SGL fused-gate API used by the GLM router.""" - - return _sgl_moe_support("MoE gate") - - def sgl_moe_alignment_support() -> tuple[bool, str]: """Check the SGL expert-alignment API used by the Triton MoE provider.""" - return _sgl_moe_support("MoE alignment") - - -def sgl_moe_ep_alignment_support() -> tuple[bool, str]: - """Check the invalid-route filtering required by expert parallelism.""" - - return _sgl_moe_support("MoE EP alignment") - - -class SglGlmFusedMoeGate: - """One-kernel sigmoid, biased top-k, normalization, and route scaling.""" - - def __init__(self, *, num_experts: int, top_k: int) -> None: - supported, reason = sgl_fused_moe_gate_support() - if not supported: - raise RuntimeError(reason) - if (int(num_experts), int(top_k)) != (64, 4): - raise ValueError( - "SGL GLM fused gate is validated only for 64 experts and " - f"top-k 4, got {num_experts} and {top_k}." - ) - installed = _installed_sgl_kernel_version() - if installed is None: - raise RuntimeError("sgl-kernel package metadata is unavailable") - _, parsed = installed - self._api_series = parsed[:2] - if self._api_series == (0, 3): - from sgl_kernel import moe_fused_gate - - self._op = moe_fused_gate - else: - from sgl_kernel import topk_sigmoid - - self._op = topk_sigmoid - - def __call__( - self, - router_logits: torch.Tensor, - correction_bias: torch.Tensor, - *, - top_k: int, - routed_scaling_factor: float, - ) -> tuple[torch.Tensor, torch.Tensor]: - if tuple(router_logits.shape[1:]) != (64,) or int(top_k) != 4: - raise ValueError( - "SGL GLM fused gate expects logits [tokens, 64] and top-k 4, " - f"got {tuple(router_logits.shape)} and {top_k}." - ) - if router_logits.dtype != torch.float32: - raise TypeError( - "SGL GLM fused gate requires FP32 logits, got " - f"{router_logits.dtype}." - ) - if tuple(correction_bias.shape) != (64,): - raise ValueError( - "SGL GLM fused gate expects correction_bias [64], got " - f"{tuple(correction_bias.shape)}." - ) - - # The 0.3.x kernel supports at most 32 experts per group. Splitting - # GLM's single 64-expert group into two virtual groups and selecting - # both preserves the original global top-k candidate set. - if self._api_series == (0, 3): - weights, ids = self._op( - router_logits, - correction_bias, - 2, - 2, - 4, - 0, - float(routed_scaling_factor), - True, - ) - else: - weights = torch.empty( - (int(router_logits.shape[0]), 4), - dtype=torch.float32, - device=router_logits.device, - ) - ids = torch.empty( - (int(router_logits.shape[0]), 4), - dtype=torch.int32, - device=router_logits.device, - ) - self._op(weights, ids, router_logits, True, correction_bias) - weights.mul_(float(routed_scaling_factor)) - return weights, ids + return sgl_kernel_support("MoE alignment") def sgl_moe_align_block_size( @@ -187,27 +34,19 @@ def sgl_moe_align_block_size( supported, reason = sgl_moe_alignment_support() if not supported: raise RuntimeError(reason) - installed = _installed_sgl_kernel_version() - if installed is None: - raise RuntimeError("sgl-kernel package metadata is unavailable") - _, parsed = installed - local_topk_ids = topk_ids has_remote_experts = num_local_experts != int(num_experts) - ignore_invalid_expert = has_remote_experts and parsed[:2] == (0, 4) - if num_local_experts != int(num_experts): + if has_remote_experts: from sparsevllm.triton_kernel.moe import localize_expert_ids local_topk_ids = localize_expert_ids( topk_ids, local_expert_start=local_expert_start, local_expert_end=local_expert_end, - remote_expert_id=-1 if ignore_invalid_expert else num_local_experts, + remote_expert_id=num_local_experts, ) num_assignments = int(local_topk_ids.numel()) - padding_experts = num_local_experts + int( - has_remote_experts and parsed[:2] == (0, 3) - ) + padding_experts = num_local_experts + int(has_remote_experts) max_num_tokens_padded = triton.cdiv( num_assignments + padding_experts * (int(block_size) - 1), int(block_size), @@ -234,31 +73,18 @@ def sgl_moe_align_block_size( ) from sgl_kernel import moe_align_block_size - # Both validated APIs iterate to num_experts - 1. The extra empty logical + # sgl-kernel 0.3.x iterates to num_experts - 1. The extra empty logical # expert makes the complete [0, num_experts) range participate. - if parsed[:2] == (0, 3): - moe_align_block_size( - local_topk_ids, - num_local_experts + 1, - int(block_size), - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - cumsum_buffer, - True, - ) - else: - moe_align_block_size( - local_topk_ids, - num_local_experts + 1, - int(block_size), - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - cumsum_buffer, - True, - ignore_invalid_expert, - ) + moe_align_block_size( + local_topk_ids, + num_local_experts + 1, + int(block_size), + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + cumsum_buffer, + True, + ) return MoeAlignment( sorted_token_ids=sorted_token_ids, expert_ids=expert_ids, @@ -269,9 +95,6 @@ def sgl_moe_align_block_size( __all__ = [ - "SglGlmFusedMoeGate", - "sgl_fused_moe_gate_support", "sgl_moe_alignment_support", - "sgl_moe_ep_alignment_support", "sgl_moe_align_block_size", ] diff --git a/tests/test_glm_cuda_graph.py b/tests/test_glm_cuda_graph.py index 85ac0f5e..6dda45da 100644 --- a/tests/test_glm_cuda_graph.py +++ b/tests/test_glm_cuda_graph.py @@ -12,7 +12,7 @@ from torch import nn from sparsevllm.config import RuntimeLayout -from sparsevllm.configs.model import resolve_attention_qk_head_dim +from sparsevllm.models.layout import resolve_attention_qk_head_dim from sparsevllm.distributed import ParallelContext from sparsevllm.engine.cache_manager import LayerBatchStates from sparsevllm.engine.cache_manager.h2o import H2OCacheManager diff --git a/tests/test_glm_runtime_compatibility.py b/tests/test_glm_runtime_compatibility.py index d11a4f82..db748c3a 100644 --- a/tests/test_glm_runtime_compatibility.py +++ b/tests/test_glm_runtime_compatibility.py @@ -279,19 +279,3 @@ def test_glm_config_accepts_sparse_latent_layout(method): assert config.attention_cache_layout == CacheLayout.MLA_LATENT.value assert config.vllm_sparse_method == method - - -@pytest.mark.parametrize( - ("hf_override", "message"), - [ - ({"torch_dtype": torch.float16}, "requires BF16"), - ({"kv_lora_rank": 256}, "requires kv_lora_rank=512"), - ({"qk_rope_head_dim": 32}, "qk_rope_head_dim=64"), - ], -) -def test_glm_config_rejects_unsupported_mla_storage_contract( - hf_override, - message, -): - with pytest.raises(NotImplementedError, match=message): - _glm_config(hf_overrides=hf_override) diff --git a/tests/test_operator_providers.py b/tests/test_operator_providers.py index 808ad431..9928e7ed 100644 --- a/tests/test_operator_providers.py +++ b/tests/test_operator_providers.py @@ -571,7 +571,7 @@ def test_glm_tp2_ep2_moe_uses_sgl_aligned_provider(): ) with patch( - "sparsevllm.operators.sgl_moe.sgl_moe_ep_alignment_support", + "sparsevllm.operators.sgl_moe.sgl_moe_alignment_support", return_value=(True, "supported"), ): resolved = OpResolver(MOE_REGISTRY).resolve( diff --git a/tests/test_qwen35_mixed_runtime.py b/tests/test_qwen35_mixed_runtime.py index 75a54a29..27cc1fc8 100644 --- a/tests/test_qwen35_mixed_runtime.py +++ b/tests/test_qwen35_mixed_runtime.py @@ -907,6 +907,10 @@ def stop_at_model_construction(_config): parallel_topology=ParallelTopology(1, 1, 1), uses_outer_tp_moe_layout=False, mlp_chunk_size=16384, + decode_cuda_graph=False, + decode_cuda_graph_capture_sizes=None, + max_decoding_seqs=64, + max_num_seqs_in_batch=32, hf_config=SimpleNamespace(model_type="qwen2", torch_dtype=torch.float32), model_spec=resolve_model_spec("qwen2"), ) diff --git a/tests/test_sgl_fa3.py b/tests/test_sgl_fa3.py index db576d95..cbb7fb28 100644 --- a/tests/test_sgl_fa3.py +++ b/tests/test_sgl_fa3.py @@ -1,6 +1,5 @@ from __future__ import annotations -import importlib.metadata from unittest.mock import patch import pytest @@ -8,15 +7,10 @@ from sparsevllm.operators.sgl_fa3 import ( SglFa3DecodeKernel, - _parse_version, sgl_fa3_support, ) -def test_sgl_fa3_version_parser_accepts_post_release() -> None: - assert _parse_version("0.3.14.post1+cu128") == (0, 3, 14) - - def test_sgl_fa3_support_rejects_missing_package() -> None: with patch("importlib.util.find_spec", return_value=None): assert sgl_fa3_support() == (False, "sgl-kernel is not installed") @@ -25,39 +19,34 @@ def test_sgl_fa3_support_rejects_missing_package() -> None: def test_sgl_fa3_support_rejects_old_package() -> None: with ( patch("importlib.util.find_spec", return_value=object()), - patch("importlib.metadata.version", return_value="0.3.13"), + patch("importlib.metadata.version", return_value="0.3.20"), ): supported, reason = sgl_fa3_support() assert not supported - assert ">= 0.3.14" in reason + assert ">=0.3.21,<0.4" in reason -def test_sgl_fa3_support_rejects_old_04_package() -> None: +def test_sgl_fa3_support_rejects_04_package() -> None: with ( patch("importlib.util.find_spec", return_value=object()), - patch("importlib.metadata.version", return_value="0.4.0"), + patch("importlib.metadata.version", return_value="0.4.5"), ): supported, reason = sgl_fa3_support() assert not supported - assert "sglang-kernel >= 0.4.5" in reason - + assert ">=0.3.21,<0.4" in reason -def test_sgl_fa3_support_accepts_sglang_kernel_04_package() -> None: - def package_version(distribution: str) -> str: - if distribution == "sglang-kernel": - return "0.4.5" - raise importlib.metadata.PackageNotFoundError(distribution) +def test_sgl_fa3_support_accepts_declared_package() -> None: with ( patch("importlib.util.find_spec", return_value=object()), - patch("importlib.metadata.version", side_effect=package_version), + patch("importlib.metadata.version", return_value="0.3.21"), ): supported, reason = sgl_fa3_support() assert supported - assert "0.4.5" in reason + assert "0.3.21" in reason @pytest.mark.skipif( diff --git a/tests/test_sgl_moe.py b/tests/test_sgl_moe.py index d318fad4..f60c29c9 100644 --- a/tests/test_sgl_moe.py +++ b/tests/test_sgl_moe.py @@ -1,133 +1,50 @@ from __future__ import annotations -import importlib.metadata from unittest.mock import patch import pytest import torch from sparsevllm.operators.sgl_moe import ( - SglGlmFusedMoeGate, - _parse_version, - sgl_fused_moe_gate_support, sgl_moe_alignment_support, sgl_moe_align_block_size, - sgl_moe_ep_alignment_support, ) from sparsevllm.triton_kernel.moe import moe_align_block_size from sparsevllm.triton_kernel.moe import fused_moe -def test_sgl_moe_version_parser_accepts_post_release() -> None: - assert _parse_version("0.3.14.post1+cu128") == (0, 3, 14) - - def test_sgl_moe_support_rejects_missing_package() -> None: with patch("importlib.util.find_spec", return_value=None): - assert sgl_fused_moe_gate_support() == ( + assert sgl_moe_alignment_support() == ( False, "sgl-kernel is not installed", ) -def test_sgl_moe_support_accepts_sglang_kernel_04_package() -> None: - def package_version(distribution: str) -> str: - if distribution == "sglang-kernel": - return "0.4.5" - raise importlib.metadata.PackageNotFoundError(distribution) - +def test_sgl_moe_support_rejects_04_package() -> None: with ( patch("importlib.util.find_spec", return_value=object()), - patch("importlib.metadata.version", side_effect=package_version), + patch("importlib.metadata.version", return_value="0.4.5"), ): - gate_supported, gate_reason = sgl_fused_moe_gate_support() - align_supported, align_reason = sgl_moe_alignment_support() - ep_supported, ep_reason = sgl_moe_ep_alignment_support() + supported, reason = sgl_moe_alignment_support() - assert gate_supported - assert align_supported - assert ep_supported - assert "0.4.5" in gate_reason - assert "0.4.5" in align_reason - assert "0.4.5" in ep_reason + assert not supported + assert ">=0.3.21,<0.4" in reason -def test_sgl_moe_ep_support_accepts_canonical_03_package() -> None: - def package_version(distribution: str) -> str: - if distribution == "sgl-kernel": - return "0.3.21" - raise importlib.metadata.PackageNotFoundError(distribution) - +def test_sgl_moe_support_accepts_declared_package() -> None: with ( patch("importlib.util.find_spec", return_value=object()), - patch("importlib.metadata.version", side_effect=package_version), + patch("importlib.metadata.version", return_value="0.3.21"), ): - supported, reason = sgl_moe_ep_alignment_support() + supported, reason = sgl_moe_alignment_support() assert supported assert "0.3.21" in reason @pytest.mark.skipif( - not torch.cuda.is_available() or not sgl_fused_moe_gate_support()[0], - reason="CUDA and a validated sgl-kernel are required", -) -def test_sgl_glm_fused_gate_matches_reference_and_replays_graph() -> None: - torch.manual_seed(20260807) - logits = torch.randn(32, 64, dtype=torch.float32, device="cuda") * 3 - correction_bias = ( - torch.randn(64, dtype=torch.float32, device="cuda") * 0.1 - ) - scaling = 1.8 - scores = torch.sigmoid(logits) - reference_ids = torch.topk( - scores + correction_bias, - 4, - dim=-1, - sorted=False, - ).indices - reference_weights = scores.gather(1, reference_ids) - reference_weights /= reference_weights.sum(dim=-1, keepdim=True) + 1e-20 - reference_weights *= scaling - gate = SglGlmFusedMoeGate(num_experts=64, top_k=4) - - actual_weights, actual_ids = gate( - logits, - correction_bias, - top_k=4, - routed_scaling_factor=scaling, - ) - - reference_order = reference_ids.argsort(dim=-1) - actual_order = actual_ids.argsort(dim=-1) - sorted_reference_ids = reference_ids.gather(1, reference_order) - sorted_actual_ids = actual_ids.gather(1, actual_order) - sorted_reference_weights = reference_weights.gather(1, reference_order) - sorted_actual_weights = actual_weights.gather(1, actual_order) - assert torch.equal(sorted_actual_ids, sorted_reference_ids) - torch.testing.assert_close( - sorted_actual_weights, - sorted_reference_weights, - rtol=2e-6, - atol=2e-7, - ) - - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - graph_weights, graph_ids = gate( - logits, - correction_bias, - top_k=4, - routed_scaling_factor=scaling, - ) - graph.replay() - torch.cuda.synchronize() - assert torch.equal(graph_ids, actual_ids) - torch.testing.assert_close(graph_weights, actual_weights, rtol=0, atol=0) - - -@pytest.mark.skipif( - not torch.cuda.is_available() or not sgl_fused_moe_gate_support()[0], + not torch.cuda.is_available() or not sgl_moe_alignment_support()[0], reason="CUDA and a validated sgl-kernel are required", ) def test_sgl_moe_alignment_matches_grouped_assignments() -> None: @@ -176,7 +93,7 @@ def test_sgl_moe_alignment_matches_grouped_assignments() -> None: @pytest.mark.skipif( - not torch.cuda.is_available() or not sgl_moe_ep_alignment_support()[0], + not torch.cuda.is_available() or not sgl_moe_alignment_support()[0], reason="CUDA and a validated sgl-kernel are required", ) @pytest.mark.parametrize(("local_start", "local_end"), [(0, 32), (32, 64)]) @@ -239,7 +156,7 @@ def grouped_assignments(sorted_ids, expert_ids, count): @pytest.mark.skipif( - not torch.cuda.is_available() or not sgl_moe_ep_alignment_support()[0], + not torch.cuda.is_available() or not sgl_moe_alignment_support()[0], reason="CUDA and a validated sgl-kernel are required", ) @pytest.mark.parametrize("num_tokens", [1, 2, 4, 5, 64]) From 5645405f31f08bab4b722bc3a2af3258ec049947 Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Wed, 12 Aug 2026 16:13:15 +0800 Subject: [PATCH 26/48] test: remove glm validation artifacts --- .../compare_glm_parallel_variants.py | 293 --------- .../validate_glm_checkpoint_parallel.py | 558 ------------------ tests/test_compare_glm_parallel_variants.py | 104 ---- .../test_validate_glm_checkpoint_parallel.py | 232 -------- 4 files changed, 1187 deletions(-) delete mode 100644 scripts/validation/compare_glm_parallel_variants.py delete mode 100644 scripts/validation/validate_glm_checkpoint_parallel.py delete mode 100644 tests/test_compare_glm_parallel_variants.py delete mode 100644 tests/test_validate_glm_checkpoint_parallel.py diff --git a/scripts/validation/compare_glm_parallel_variants.py b/scripts/validation/compare_glm_parallel_variants.py deleted file mode 100644 index 9d62eca2..00000000 --- a/scripts/validation/compare_glm_parallel_variants.py +++ /dev/null @@ -1,293 +0,0 @@ -#!/usr/bin/env python3 -"""Compare two GLM checkpoint parallel-validation artifacts.""" - -from __future__ import annotations - -import argparse -import hashlib -import json -import sys -import traceback -from pathlib import Path -from typing import Any - -import torch - - -LOGIT_KEYS = ("prefill", "final_decode") -MATCHED_RUN_FIELDS = ( - "model_path", - "seed", - "prompt_len", - "prompt_sha256", - "max_tokens", - "sparse_method", - "tensor_parallel_size", -) - - -def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser( - description="Compare real-checkpoint GLM parallel validation artifacts." - ) - parser.add_argument("--reference_dir", required=True) - parser.add_argument("--candidate_dir", required=True) - parser.add_argument("--output", required=True) - parser.add_argument("--max_abs_diff", type=float, required=True) - parser.add_argument("--mean_abs_diff", type=float, required=True) - parser.add_argument("--p99_abs_diff", type=float, required=True) - parser.add_argument( - "--allow_first_token_tie", - action=argparse.BooleanOptionalAction, - default=False, - help=( - "Accept a first generated-token mismatch when the candidate token " - "is tied for the maximum in the reference prefill logits and all " - "later generated tokens match." - ), - ) - parser.add_argument( - "--tie_atol", - type=float, - default=0.0, - help="Absolute tolerance for the reference-logit tie check.", - ) - return parser.parse_args() - - -def _load_json(path: Path) -> dict[str, Any]: - payload = json.loads(path.read_text(encoding="utf-8")) - if not isinstance(payload, dict): - raise TypeError(f"Expected a JSON object in {path}.") - return payload - - -def _file_sha256(path: Path) -> str: - digest = hashlib.sha256() - with path.open("rb") as handle: - for chunk in iter(lambda: handle.read(1024 * 1024), b""): - digest.update(chunk) - return digest.hexdigest() - - -def _load_artifact(run_dir: Path) -> tuple[dict[str, Any], dict[str, torch.Tensor]]: - summary_path = run_dir / "summary.json" - logits_path = run_dir / "logits.pt" - if not summary_path.is_file() or not logits_path.is_file(): - raise FileNotFoundError( - f"Expected summary.json and logits.pt under {run_dir}." - ) - summary = _load_json(summary_path) - if summary.get("status") != "success": - raise RuntimeError( - f"Parallel validation did not succeed in {run_dir}: " - f"status={summary.get('status')!r}." - ) - raw_logits = torch.load(logits_path, map_location="cpu", weights_only=True) - if not isinstance(raw_logits, dict): - raise TypeError(f"Expected a tensor dictionary in {logits_path}.") - logits: dict[str, torch.Tensor] = {} - for key in LOGIT_KEYS: - tensor = raw_logits.get(key) - if not isinstance(tensor, torch.Tensor): - raise TypeError(f"Missing tensor {key!r} in {logits_path}.") - logits[key] = tensor.detach().float().cpu().contiguous() - return summary, logits - - -def _compare_tensor( - reference: torch.Tensor, - candidate: torch.Tensor, - *, - max_abs_limit: float, - mean_abs_limit: float, - p99_abs_limit: float, -) -> dict[str, Any]: - if reference.shape != candidate.shape: - raise ValueError( - "Logit shape mismatch: " - f"reference={tuple(reference.shape)} candidate={tuple(candidate.shape)}." - ) - if not torch.isfinite(reference).all() or not torch.isfinite(candidate).all(): - raise ValueError("Logit artifacts must contain only finite values.") - diff = (reference - candidate).abs().flatten() - max_abs = float(diff.max().item()) - mean_abs = float(diff.mean().item()) - p99_abs = float(torch.quantile(diff, 0.99).item()) - reference_argmax = reference.argmax(dim=-1).tolist() - candidate_argmax = candidate.argmax(dim=-1).tolist() - gates = { - "max_abs_diff": max_abs <= float(max_abs_limit), - "mean_abs_diff": mean_abs <= float(mean_abs_limit), - "p99_abs_diff": p99_abs <= float(p99_abs_limit), - } - return { - "shape": list(reference.shape), - "max_abs_diff": max_abs, - "mean_abs_diff": mean_abs, - "p99_abs_diff": p99_abs, - "reference_argmax": reference_argmax, - "candidate_argmax": candidate_argmax, - "argmax_match": reference_argmax == candidate_argmax, - "limits": { - "max_abs_diff": float(max_abs_limit), - "mean_abs_diff": float(mean_abs_limit), - "p99_abs_diff": float(p99_abs_limit), - }, - "gates": gates, - "passed": all(gates.values()), - } - - -def _validate_generated_tokens( - reference_summary: dict[str, Any], - candidate_summary: dict[str, Any], - reference_prefill_logits: torch.Tensor, - *, - allow_first_token_tie: bool, - tie_atol: float, -) -> dict[str, Any]: - reference = [int(value) for value in reference_summary["generated_token_ids"]] - candidate = [int(value) for value in candidate_summary["generated_token_ids"]] - exact_match = reference == candidate - tie_accepted = False - reference_tie_gap = None - if ( - not exact_match - and allow_first_token_tie - and reference - and len(reference) == len(candidate) - and reference[1:] == candidate[1:] - ): - row = reference_prefill_logits.reshape( - -1, reference_prefill_logits.shape[-1] - )[-1] - candidate_first = candidate[0] - if not 0 <= candidate_first < int(row.numel()): - raise ValueError( - f"Candidate token {candidate_first} is outside the logit vocabulary." - ) - reference_tie_gap = float((row.max() - row[candidate_first]).item()) - tie_accepted = reference_tie_gap <= float(tie_atol) - return { - "reference": reference, - "candidate": candidate, - "exact_match": exact_match, - "first_token_tie_accepted": tie_accepted, - "reference_tie_gap": reference_tie_gap, - "tie_atol": float(tie_atol), - "passed": exact_match or tie_accepted, - } - - -def compare_runs( - reference_dir: Path, - candidate_dir: Path, - *, - max_abs_limit: float, - mean_abs_limit: float, - p99_abs_limit: float, - allow_first_token_tie: bool, - tie_atol: float, -) -> dict[str, Any]: - limits = (max_abs_limit, mean_abs_limit, p99_abs_limit, tie_atol) - if any(float(value) < 0.0 for value in limits): - raise ValueError(f"Comparison tolerances must be non-negative, got {limits}.") - reference_summary, reference_logits = _load_artifact(reference_dir) - candidate_summary, candidate_logits = _load_artifact(candidate_dir) - - mismatches = { - field: { - "reference": reference_summary.get(field), - "candidate": candidate_summary.get(field), - } - for field in MATCHED_RUN_FIELDS - if reference_summary.get(field) != candidate_summary.get(field) - } - if mismatches: - raise ValueError(f"Parallel runs are not comparable: {mismatches}.") - - comparisons = { - key: _compare_tensor( - reference_logits[key], - candidate_logits[key], - max_abs_limit=max_abs_limit, - mean_abs_limit=mean_abs_limit, - p99_abs_limit=p99_abs_limit, - ) - for key in LOGIT_KEYS - } - generated_tokens = _validate_generated_tokens( - reference_summary, - candidate_summary, - reference_logits["prefill"], - allow_first_token_tie=allow_first_token_tie, - tie_atol=tie_atol, - ) - gates = { - **{f"{key}_logits": value["passed"] for key, value in comparisons.items()}, - "generated_tokens": generated_tokens["passed"], - } - return { - "status": "success" if all(gates.values()) else "metric_failed", - "reference": { - "directory": str(reference_dir.resolve()), - "tensor_parallel_size": reference_summary["tensor_parallel_size"], - "expert_parallel_size": reference_summary["expert_parallel_size"], - "summary_sha256": _file_sha256(reference_dir / "summary.json"), - "logits_sha256": _file_sha256(reference_dir / "logits.pt"), - }, - "candidate": { - "directory": str(candidate_dir.resolve()), - "tensor_parallel_size": candidate_summary["tensor_parallel_size"], - "expert_parallel_size": candidate_summary["expert_parallel_size"], - "summary_sha256": _file_sha256(candidate_dir / "summary.json"), - "logits_sha256": _file_sha256(candidate_dir / "logits.pt"), - }, - "logits": comparisons, - "generated_tokens": generated_tokens, - "gates": gates, - } - - -def main() -> int: - args = _parse_args() - output_path = Path(args.output).expanduser().resolve() - output_path.parent.mkdir(parents=True, exist_ok=True) - try: - result = compare_runs( - Path(args.reference_dir).expanduser().resolve(), - Path(args.candidate_dir).expanduser().resolve(), - max_abs_limit=float(args.max_abs_diff), - mean_abs_limit=float(args.mean_abs_diff), - p99_abs_limit=float(args.p99_abs_diff), - allow_first_token_tie=bool(args.allow_first_token_tie), - tie_atol=float(args.tie_atol), - ) - except Exception as exc: - result = { - "status": "invalid_input", - "reference_dir": str( - Path(args.reference_dir).expanduser().resolve() - ), - "candidate_dir": str( - Path(args.candidate_dir).expanduser().resolve() - ), - "error_type": type(exc).__name__, - "error_message": str(exc), - "traceback": traceback.format_exc(), - } - output_path.write_text( - json.dumps(result, ensure_ascii=False, indent=2, sort_keys=True) + "\n", - encoding="utf-8", - ) - stream = sys.stdout if result["status"] == "success" else sys.stderr - print( - json.dumps(result, ensure_ascii=False, indent=2, sort_keys=True), - file=stream, - ) - return 0 if result["status"] == "success" else 1 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/scripts/validation/validate_glm_checkpoint_parallel.py b/scripts/validation/validate_glm_checkpoint_parallel.py deleted file mode 100644 index c2b5b5f2..00000000 --- a/scripts/validation/validate_glm_checkpoint_parallel.py +++ /dev/null @@ -1,558 +0,0 @@ -#!/usr/bin/env python3 -"""Validate one real GLM-4.7-Flash TP/EP runtime variant. - -This script deliberately uses the public ``LLM`` entry point so worker launch, -weight loading, scheduling, collectives, and cache allocation follow the same -path as serving. Cross-variant comparisons are performed separately from the -saved float32 logits artifacts. -""" - -from __future__ import annotations - -import argparse -import hashlib -import json -import os -import subprocess -import sys -import traceback -from datetime import datetime -from pathlib import Path -from typing import Any - -import torch - - -REPO_ROOT = Path(__file__).resolve().parents[2] - - -def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser( - description="Validate one real GLM checkpoint TP/EP variant." - ) - parser.add_argument("--model_path", required=True) - parser.add_argument("--output_dir", required=True) - parser.add_argument("--tensor_parallel_size", type=int, default=1) - parser.add_argument("--expert_parallel_size", type=int, default=1) - parser.add_argument( - "--sparse_method", - choices=( - "vanilla", - "streamingllm", - "snapkv", - "h2o", - "omnikv", - "rkv", - ), - default="vanilla", - ) - parser.add_argument( - "--decode_cuda_graph", - action=argparse.BooleanOptionalAction, - default=False, - ) - parser.add_argument("--seed", type=int, default=20260805) - parser.add_argument("--prompt_len", type=int, default=64) - parser.add_argument("--max_tokens", type=int, default=3) - parser.add_argument("--gpu_memory_utilization", type=float, default=0.80) - parser.add_argument("--mlp_chunk_size", type=int, default=64) - parser.add_argument("--max_steps", type=int, default=256) - return parser.parse_args() - - -def _write_json(path: Path, payload: Any) -> None: - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text( - json.dumps(payload, ensure_ascii=False, indent=2, sort_keys=True) + "\n", - encoding="utf-8", - ) - - -def _git_output(*args: str) -> str: - command = ["git", *args] - try: - result = subprocess.run( - command, - cwd=REPO_ROOT, - check=True, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - text=True, - ) - except (FileNotFoundError, subprocess.CalledProcessError) as exc: - detail = getattr(exc, "stderr", None) or str(exc) - raise RuntimeError( - f"Git provenance command failed in {REPO_ROOT}: " - f"{' '.join(command)}: {str(detail).strip()}" - ) from exc - return result.stdout.strip() - - -def _tensor_sha256(tensor: torch.Tensor) -> str: - value = tensor.detach().float().cpu().contiguous() - return hashlib.sha256(value.numpy().tobytes()).hexdigest() - - -def _tensor_evidence(tensor: torch.Tensor) -> dict[str, Any]: - value = tensor.detach().float().cpu().contiguous() - return { - "shape": [int(dim) for dim in value.shape], - "dtype": str(value.dtype), - "sha256": _tensor_sha256(value), - "argmax": value.argmax(dim=-1).tolist(), - "mean": float(value.mean().item()), - "max": float(value.max().item()), - "min": float(value.min().item()), - } - - -def _prompt_token_ids(*, seed: int, prompt_len: int, vocab_size: int) -> list[int]: - if prompt_len <= 0: - raise ValueError(f"prompt_len must be positive, got {prompt_len}.") - if vocab_size <= 1024: - raise ValueError( - "GLM checkpoint validation expects a vocabulary larger than 1024, " - f"got {vocab_size}." - ) - usable = vocab_size - 512 - base = 256 + int(seed) % max(1, usable - 256) - return [256 + ((base + 97 * index) % usable) for index in range(prompt_len)] - - -def _validate_parallel_evidence( - summaries: list[dict[str, Any]], - *, - tensor_parallel_size: int, - expert_parallel_size: int, - last_moe_layer: int, - require_all_ep_ranks_hit: bool = True, - require_decode_cuda_graph: bool = False, - require_decode_cuda_graph_execution: bool = False, -) -> dict[str, Any]: - from sparsevllm.distributed.parallel_context import ( - hybrid_moe_group_ranks, - parallel_group_ranks, - ) - - tp_size = int(tensor_parallel_size) - ep_size = int(expert_parallel_size) - hybrid_moe = tp_size > 1 and ep_size > 1 - expected_world_size = tp_size if hybrid_moe else tp_size * ep_size - if len(summaries) != expected_world_size: - raise RuntimeError( - "Expected one debug summary per world rank: " - f"expected={expected_world_size} got={len(summaries)}." - ) - - world_ranks = sorted(int(item["world_rank"]) for item in summaries) - if world_ranks != list(range(expected_world_size)): - raise RuntimeError(f"Unexpected world ranks: {world_ranks}.") - - if hybrid_moe: - raw_groups = hybrid_moe_group_ranks( - outer_tp_size=tp_size, - moe_ep_size=ep_size, - ) - group_ranks = { - "attention": raw_groups["attention"], - "expert": raw_groups["moe_expert"], - "moe_tensor": raw_groups["moe_tensor"], - "data": raw_groups["data"], - } - else: - raw_groups = parallel_group_ranks( - tp_size=tp_size, - ep_size=ep_size, - dp_size=1, - ) - group_ranks = { - "attention": raw_groups["tensor"], - "expert": raw_groups["expert"], - "moe_tensor": raw_groups["tensor"], - "data": raw_groups["data"], - } - - def expected_group(dimension: str, world_rank: int) -> dict[str, Any]: - ranks = next( - ranks - for ranks in group_ranks[dimension] - if int(world_rank) in ranks - ) - return { - "rank": ranks.index(int(world_rank)), - "size": len(ranks), - "ranks": list(ranks), - } - - local_ranges: list[tuple[int, int]] = [] - local_hit_counts: list[int] = [] - replica_checks: list[dict[str, Any]] = [] - layer_key = str(int(last_moe_layer)) - for summary in summaries: - configured = summary["parallel"]["configured"] - expected_configured = { - "tensor_parallel_size": int(tensor_parallel_size), - "expert_parallel_size": int(expert_parallel_size), - "data_parallel_size": 1, - "world_size": expected_world_size, - } - if configured != expected_configured: - raise RuntimeError( - "Configured parallel evidence mismatch: " - f"expected={expected_configured} got={configured}." - ) - effective = summary["parallel"]["effective"] - world = effective["world"] - if int(world["size"]) != expected_world_size: - raise RuntimeError(f"Effective world-size mismatch: {world}.") - if world != { - "rank": int(summary["world_rank"]), - "size": expected_world_size, - "ranks": list(range(expected_world_size)), - }: - raise RuntimeError( - "Effective world group mismatch: " - f"world_rank={summary['world_rank']} group={world}." - ) - for dimension in ("attention", "expert", "moe_tensor", "data"): - expected = expected_group(dimension, int(summary["world_rank"])) - if effective[dimension] != expected: - raise RuntimeError( - f"Effective {dimension} group mismatch on world rank " - f"{summary['world_rank']}: expected={expected} " - f"got={effective[dimension]}." - ) - - layer = summary.get("moe_local", {}).get(layer_key) - if layer is None: - raise RuntimeError( - f"Missing MoE debug evidence for layer={last_moe_layer} " - f"on world_rank={summary['world_rank']}." - ) - local_ranges.append( - (int(layer["local_expert_start"]), int(layer["local_expert_end"])) - ) - local_hit_counts.append(int(layer["local_hit_count"])) - - consistency = summary.get("replica_consistency") - if not isinstance(consistency, dict): - raise RuntimeError( - "Missing replica-consistency evidence for " - f"world_rank={summary['world_rank']}." - ) - replica_checks.append(consistency) - logits_tolerance_ratio = consistency["last_logits_tolerance_ratio"] - if tensor_parallel_size > 1: - if ( - logits_tolerance_ratio is not None - or consistency.get("last_logits_comparison") - != "not_applicable_tp_vocab_sharded" - ): - raise RuntimeError( - "TP-sharded logits must be marked not applicable for cross-rank " - f"comparison: world_rank={summary['world_rank']} " - f"evidence={consistency}." - ) - elif logits_tolerance_ratio is None or float(logits_tolerance_ratio) > 1.0: - raise RuntimeError( - "Cross-rank logits exceeded the declared tolerance or were not " - f"compared: world_rank={summary['world_rank']} evidence={consistency}." - ) - moe_consistency = consistency.get("moe_layers", {}).get(layer_key) - if moe_consistency is None: - raise RuntimeError( - f"Missing cross-rank MoE consistency for layer={layer_key}." - ) - if bool(moe_consistency["topk_ids_mismatch"]): - raise RuntimeError( - f"Router top-k differs across ranks: {moe_consistency}." - ) - for field in ("topk_weights_tolerance_ratio", "output_tolerance_ratio"): - if float(moe_consistency[field]) > 1.0: - raise RuntimeError( - f"Cross-rank MoE {field} exceeded tolerance: " - f"{moe_consistency}." - ) - graph = summary.get("decode_cuda_graph") - if not isinstance(graph, dict): - raise RuntimeError( - "Missing decode CUDA Graph evidence on world rank " - f"{summary['world_rank']}." - ) - if bool(graph.get("enabled")) != bool(require_decode_cuda_graph): - raise RuntimeError( - "Decode CUDA Graph config evidence mismatch on world rank " - f"{summary['world_rank']}: {graph}." - ) - if require_decode_cuda_graph_execution: - if not require_decode_cuda_graph: - raise ValueError( - "Requiring decode CUDA Graph execution also requires the " - "graph config to be enabled." - ) - if int(graph.get("capture_count", 0)) <= 0: - raise RuntimeError( - f"World rank {summary['world_rank']} did not capture a decode graph." - ) - if int(graph.get("replay_count", 0)) <= 0: - raise RuntimeError( - f"World rank {summary['world_rank']} did not replay a decode graph." - ) - if int(graph.get("eager_static_count", 0)) != 0 or int( - graph.get("force_eager_count", 0) - ) != 0: - raise RuntimeError( - "Decode CUDA Graph used an eager path on world rank " - f"{summary['world_rank']}: {graph}." - ) - - if ep_size > 1: - expected_width = 64 // ep_size - expected_ranges = [ - (rank * expected_width, (rank + 1) * expected_width) - for rank in range(ep_size) - for _ in range(tp_size // ep_size if hybrid_moe else tp_size) - ] - if sorted(local_ranges) != expected_ranges: - raise RuntimeError( - "EP local expert shards do not partition 64 experts: " - f"expected={expected_ranges} got={sorted(local_ranges)}." - ) - hits_by_range: dict[tuple[int, int], int] = {} - for local_range, hit_count in zip(local_ranges, local_hit_counts): - hits_by_range[local_range] = ( - hits_by_range.get(local_range, 0) + int(hit_count) - ) - if require_all_ep_ranks_hit and any( - hit_count <= 0 for hit_count in hits_by_range.values() - ): - raise RuntimeError( - "At least one EP rank did not execute a local routed expert: " - f"hits_by_range={hits_by_range}." - ) - expected_replicated_attention = tp_size == 1 - if any( - bool(summary["parallel"]["attention_replicated_for_ep"]) - != expected_replicated_attention - for summary in summaries - ): - raise RuntimeError( - "EP replicated-attention evidence does not match the selected " - f"layout: TP={tp_size}, EP={ep_size}." - ) - - return { - "world_ranks": world_ranks, - "local_expert_ranges": [list(value) for value in local_ranges], - "local_expert_hit_counts": local_hit_counts, - "all_ep_ranks_hit": all(hit_count > 0 for hit_count in local_hit_counts), - "hybrid_moe": hybrid_moe, - "expected_groups": { - key: [list(ranks) for ranks in value] - for key, value in group_ranks.items() - }, - "replica_checks": replica_checks, - } - - -def _run(args: argparse.Namespace, output_dir: Path) -> dict[str, Any]: - os.environ["SPARSEVLLM_DEBUG_RUNTIME"] = "1" - os.environ["SPARSEVLLM_DEBUG_MOE"] = "1" - - from sparsevllm import LLM, SamplingParams - - tp_size = int(args.tensor_parallel_size) - ep_size = int(args.expert_parallel_size) - if tp_size <= 0 or ep_size <= 0: - raise ValueError(f"TP and EP must be positive, got TP={tp_size}, EP={ep_size}.") - if tp_size > 1 and ep_size > 1 and tp_size % ep_size: - raise ValueError( - "Joint GLM outer TP and MoE EP requires TP divisible by EP, got " - f"TP={tp_size}, EP={ep_size}." - ) - expected_devices = ( - tp_size if tp_size > 1 and ep_size > 1 else tp_size * ep_size - ) - actual_devices = torch.cuda.device_count() - if actual_devices != expected_devices: - raise RuntimeError( - "CUDA_VISIBLE_DEVICES must expose exactly the requested world size: " - f"expected={expected_devices} got={actual_devices}." - ) - - torch.manual_seed(int(args.seed)) - max_model_len = int(args.prompt_len) + int(args.max_tokens) + 16 - llm = None - try: - llm = LLM( - args.model_path, - tensor_parallel_size=tp_size, - expert_parallel_size=ep_size, - data_parallel_size=1, - sparse_method=args.sparse_method, - enforce_eager=True, - decode_cuda_graph=bool(args.decode_cuda_graph), - max_model_len=max_model_len, - max_num_seqs_in_batch=1, - max_decoding_seqs=1, - max_num_batched_tokens=max(int(args.prompt_len), 1), - engine_prefill_chunk_size=max(int(args.prompt_len), 1), - mlp_chunk_size=int(args.mlp_chunk_size), - gpu_memory_utilization=float(args.gpu_memory_utilization), - throughput_log_interval_s=0.0, - weight_loading_workers=1, - ) - prompt = _prompt_token_ids( - seed=int(args.seed), - prompt_len=int(args.prompt_len), - vocab_size=int(llm.config.hf_config.vocab_size), - ) - sampling_params = SamplingParams( - temperature=0.0, - top_p=1.0, - ignore_eos=True, - max_tokens=int(args.max_tokens), - ) - llm.add_request(prompt, sampling_params) - - prefill_logits = None - prefill_rank_summaries = None - finished_outputs: list[tuple[int, list[int], Any, Any]] = [] - steps = 0 - while not llm.is_finished(): - step_outputs, num_tokens = llm.step() - steps += 1 - if steps > int(args.max_steps): - raise RuntimeError( - f"Validation exceeded max_steps={int(args.max_steps)}." - ) - if ( - num_tokens > 0 - and llm.last_step_token_outputs - and prefill_rank_summaries is None - ): - prefill_logits = llm.debug_last_logits().detach().float().cpu() - prefill_rank_summaries = llm.debug_sparse_state_summaries() - finished_outputs.extend(step_outputs) - - if prefill_logits is None or prefill_rank_summaries is None: - raise RuntimeError("The run completed without final-prefill evidence.") - if len(finished_outputs) != 1: - raise RuntimeError( - f"Expected one finished request, got {len(finished_outputs)}." - ) - generated_ids = [int(value) for value in finished_outputs[0][1]] - if len(generated_ids) != int(args.max_tokens): - raise RuntimeError( - "Greedy generation returned an unexpected token count: " - f"expected={int(args.max_tokens)} got={len(generated_ids)}." - ) - final_logits = llm.debug_last_logits().detach().float().cpu() - final_rank_summaries = llm.debug_sparse_state_summaries() - last_moe_layer = int(llm.config.hf_config.num_hidden_layers) - 1 - - prefill_parallel = _validate_parallel_evidence( - prefill_rank_summaries, - tensor_parallel_size=tp_size, - expert_parallel_size=ep_size, - last_moe_layer=last_moe_layer, - require_decode_cuda_graph=bool(args.decode_cuda_graph), - ) - final_parallel = _validate_parallel_evidence( - final_rank_summaries, - tensor_parallel_size=tp_size, - expert_parallel_size=ep_size, - last_moe_layer=last_moe_layer, - # A single decode token selects only top-k experts and does not - # have to cover every EP shard. Prefill coverage above proves all - # local expert partitions execute; keep the final per-rank counts - # as evidence without imposing an invalid per-token requirement. - require_all_ep_ranks_hit=False, - require_decode_cuda_graph=bool(args.decode_cuda_graph), - require_decode_cuda_graph_execution=bool(args.decode_cuda_graph), - ) - - torch.save( - { - "prefill": prefill_logits, - "final_decode": final_logits, - }, - output_dir / "logits.pt", - ) - _write_json(output_dir / "prefill_rank_summaries.json", prefill_rank_summaries) - _write_json(output_dir / "final_rank_summaries.json", final_rank_summaries) - - return { - "status": "success", - "real_checkpoint": True, - "model_path": str(Path(args.model_path).resolve()), - "tensor_parallel_size": tp_size, - "expert_parallel_size": ep_size, - "data_parallel_size": 1, - "world_size": expected_devices, - "sparse_method": str(args.sparse_method), - "decode_cuda_graph": bool(args.decode_cuda_graph), - "seed": int(args.seed), - "prompt_len": len(prompt), - "prompt_sha256": hashlib.sha256( - torch.tensor(prompt, dtype=torch.int64).numpy().tobytes() - ).hexdigest(), - "max_tokens": int(args.max_tokens), - "generated_token_ids": generated_ids, - "steps": steps, - "prefill_logits": _tensor_evidence(prefill_logits), - "final_decode_logits": _tensor_evidence(final_logits), - "prefill_parallel": prefill_parallel, - "final_parallel": final_parallel, - "logits_artifact": str(output_dir / "logits.pt"), - } - finally: - if llm is not None: - llm.exit() - - -def main() -> int: - args = _parse_args() - output_dir = Path(args.output_dir).expanduser().resolve() - output_dir.mkdir(parents=True, exist_ok=True) - run_info = { - "status": "running", - "started_at": datetime.now().astimezone().isoformat(timespec="seconds"), - "command": [sys.executable, *sys.argv], - "cwd": os.getcwd(), - "git_commit": _git_output("rev-parse", "HEAD"), - "git_status": _git_output("status", "--short"), - "cuda_visible_devices": os.getenv("CUDA_VISIBLE_DEVICES"), - "args": vars(args), - } - _write_json(output_dir / "run_info.json", run_info) - try: - summary = _run(args, output_dir) - except BaseException as exc: - failure = { - "status": "model_failed", - "error_type": type(exc).__name__, - "error_message": str(exc), - "traceback": traceback.format_exc(), - } - _write_json(output_dir / "summary.json", failure) - run_info.update( - status="model_failed", - finished_at=datetime.now().astimezone().isoformat(timespec="seconds"), - error=failure, - ) - _write_json(output_dir / "run_info.json", run_info) - raise - _write_json(output_dir / "summary.json", summary) - run_info.update( - status="success", - finished_at=datetime.now().astimezone().isoformat(timespec="seconds"), - summary_path=str(output_dir / "summary.json"), - ) - _write_json(output_dir / "run_info.json", run_info) - print(json.dumps(summary, ensure_ascii=False, indent=2, sort_keys=True)) - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/tests/test_compare_glm_parallel_variants.py b/tests/test_compare_glm_parallel_variants.py deleted file mode 100644 index 2bbbaa1a..00000000 --- a/tests/test_compare_glm_parallel_variants.py +++ /dev/null @@ -1,104 +0,0 @@ -from __future__ import annotations - -import json -from pathlib import Path - -import torch - -from scripts.validation.compare_glm_parallel_variants import compare_runs, main - - -def _write_run( - directory: Path, - *, - logits: torch.Tensor, - tokens: list[int], - ep_size: int, -) -> None: - directory.mkdir() - summary = { - "status": "success", - "model_path": "/checkpoint", - "seed": 7, - "prompt_len": 4, - "prompt_sha256": "prompt", - "max_tokens": len(tokens), - "sparse_method": "vanilla", - "tensor_parallel_size": 2, - "expert_parallel_size": ep_size, - "generated_token_ids": tokens, - } - (directory / "summary.json").write_text(json.dumps(summary), encoding="utf-8") - torch.save( - {"prefill": logits, "final_decode": logits}, - directory / "logits.pt", - ) - - -def test_compare_parallel_runs_accepts_bounded_drift_and_reference_tie(tmp_path): - reference = torch.tensor([[1.0, 3.0, 3.0, 0.0]]) - candidate = torch.tensor([[1.0, 2.5, 3.0, 0.0]]) - reference_dir = tmp_path / "reference" - candidate_dir = tmp_path / "candidate" - _write_run(reference_dir, logits=reference, tokens=[1, 3], ep_size=1) - _write_run(candidate_dir, logits=candidate, tokens=[2, 3], ep_size=2) - - result = compare_runs( - reference_dir, - candidate_dir, - max_abs_limit=0.5, - mean_abs_limit=0.2, - p99_abs_limit=0.5, - allow_first_token_tie=True, - tie_atol=0.0, - ) - - assert result["status"] == "success" - assert result["generated_tokens"]["first_token_tie_accepted"] is True - - -def test_compare_parallel_runs_fails_when_drift_exceeds_limit(tmp_path): - reference = torch.tensor([[1.0, 2.0]]) - candidate = torch.tensor([[1.0, 3.0]]) - reference_dir = tmp_path / "reference" - candidate_dir = tmp_path / "candidate" - _write_run(reference_dir, logits=reference, tokens=[1], ep_size=1) - _write_run(candidate_dir, logits=candidate, tokens=[1], ep_size=2) - - result = compare_runs( - reference_dir, - candidate_dir, - max_abs_limit=0.5, - mean_abs_limit=1.0, - p99_abs_limit=1.0, - allow_first_token_tie=False, - tie_atol=0.0, - ) - - assert result["status"] == "metric_failed" - assert result["gates"]["prefill_logits"] is False - - -def test_cli_writes_invalid_input_status_before_failing(tmp_path, monkeypatch): - output = tmp_path / "comparison.json" - monkeypatch.setattr( - "sys.argv", - [ - "compare_glm_parallel_variants.py", - "--reference_dir", - str(tmp_path / "missing-reference"), - "--candidate_dir", - str(tmp_path / "missing-candidate"), - "--output", - str(output), - "--max_abs_diff", - "1", - "--mean_abs_diff", - "1", - "--p99_abs_diff", - "1", - ], - ) - - assert main() == 1 - assert json.loads(output.read_text())["status"] == "invalid_input" diff --git a/tests/test_validate_glm_checkpoint_parallel.py b/tests/test_validate_glm_checkpoint_parallel.py deleted file mode 100644 index 6bc58bfe..00000000 --- a/tests/test_validate_glm_checkpoint_parallel.py +++ /dev/null @@ -1,232 +0,0 @@ -from __future__ import annotations - -import subprocess -from pathlib import Path - -import pytest - -from scripts.validation.validate_glm_checkpoint_parallel import ( - REPO_ROOT, - _git_output, - _validate_parallel_evidence, -) - - -def test_git_provenance_is_resolved_from_repository_root( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - monkeypatch.chdir(tmp_path) - - assert _git_output("rev-parse", "--show-toplevel") == str(REPO_ROOT) - - -def test_git_provenance_failure_is_explicit( - monkeypatch: pytest.MonkeyPatch, -) -> None: - def fail(*args, **kwargs): - del args, kwargs - raise subprocess.CalledProcessError( - returncode=128, - cmd=["git", "rev-parse", "HEAD"], - stderr="fatal: unavailable repository", - ) - - monkeypatch.setattr(subprocess, "run", fail) - - with pytest.raises( - RuntimeError, - match="Git provenance command failed.*fatal: unavailable repository", - ): - _git_output("rev-parse", "HEAD") - - -def _graph_summary(enabled: bool = False) -> dict: - return { - "enabled": enabled, - "capture_count": 1 if enabled else 0, - "replay_count": 2 if enabled else 0, - "eager_static_count": 0, - "force_eager_count": 0, - } - - -def _ep4_summary(world_rank: int, local_hit_count: int) -> dict: - layer_key = "46" - local_width = 16 - return { - "world_rank": world_rank, - "parallel": { - "configured": { - "tensor_parallel_size": 1, - "expert_parallel_size": 4, - "data_parallel_size": 1, - "world_size": 4, - }, - "effective": { - "world": { - "rank": world_rank, - "size": 4, - "ranks": [0, 1, 2, 3], - }, - "attention": { - "rank": 0, - "size": 1, - "ranks": [world_rank], - }, - "expert": { - "rank": world_rank, - "size": 4, - "ranks": [0, 1, 2, 3], - }, - "moe_tensor": { - "rank": 0, - "size": 1, - "ranks": [world_rank], - }, - "data": { - "rank": 0, - "size": 1, - "ranks": [world_rank], - }, - }, - "attention_replicated_for_ep": True, - }, - "moe_local": { - layer_key: { - "local_expert_start": world_rank * local_width, - "local_expert_end": (world_rank + 1) * local_width, - "local_hit_count": local_hit_count, - } - }, - "replica_consistency": { - "last_logits_comparison": "compared", - "last_logits_tolerance_ratio": 0.0, - "moe_layers": { - layer_key: { - "topk_ids_mismatch": False, - "topk_weights_tolerance_ratio": 0.0, - "output_tolerance_ratio": 0.0, - } - }, - }, - "decode_cuda_graph": _graph_summary(), - } - - -def test_ep_prefill_requires_every_local_expert_shard_to_execute(): - summaries = [ - _ep4_summary(rank, hit_count) - for rank, hit_count in enumerate([0, 80, 84, 92]) - ] - - with pytest.raises(RuntimeError, match="did not execute a local routed expert"): - _validate_parallel_evidence( - summaries, - tensor_parallel_size=1, - expert_parallel_size=4, - last_moe_layer=46, - ) - - -def test_ep_single_decode_keeps_zero_hit_as_evidence_without_rejecting_it(): - summaries = [ - _ep4_summary(rank, hit_count) - for rank, hit_count in enumerate([0, 1, 1, 2]) - ] - - evidence = _validate_parallel_evidence( - summaries, - tensor_parallel_size=1, - expert_parallel_size=4, - last_moe_layer=46, - require_all_ep_ranks_hit=False, - ) - - assert evidence["local_expert_hit_counts"] == [0, 1, 1, 2] - assert evidence["all_ep_ranks_hit"] is False - - -def _hybrid_tp4_ep2_summary(world_rank: int) -> dict: - ep_rank = 0 if world_rank < 2 else 1 - expert_ranks = [0, 2] if world_rank % 2 == 0 else [1, 3] - moe_tensor_ranks = [0, 1] if world_rank < 2 else [2, 3] - layer_key = "46" - return { - "world_rank": world_rank, - "parallel": { - "configured": { - "tensor_parallel_size": 4, - "expert_parallel_size": 2, - "data_parallel_size": 1, - "world_size": 4, - }, - "effective": { - "world": { - "rank": world_rank, - "size": 4, - "ranks": [0, 1, 2, 3], - }, - "attention": { - "rank": world_rank, - "size": 4, - "ranks": [0, 1, 2, 3], - }, - "expert": { - "rank": expert_ranks.index(world_rank), - "size": 2, - "ranks": expert_ranks, - }, - "moe_tensor": { - "rank": moe_tensor_ranks.index(world_rank), - "size": 2, - "ranks": moe_tensor_ranks, - }, - "data": { - "rank": 0, - "size": 1, - "ranks": [world_rank], - }, - }, - "attention_replicated_for_ep": False, - }, - "moe_local": { - layer_key: { - "local_expert_start": ep_rank * 32, - "local_expert_end": (ep_rank + 1) * 32, - "local_hit_count": 8, - } - }, - "replica_consistency": { - "last_logits_comparison": "not_applicable_tp_vocab_sharded", - "last_logits_tolerance_ratio": None, - "moe_layers": { - layer_key: { - "topk_ids_mismatch": False, - "topk_weights_tolerance_ratio": 0.0, - "output_tolerance_ratio": 0.0, - } - }, - }, - "decode_cuda_graph": _graph_summary(enabled=True), - } - - -def test_hybrid_tp_ep_validation_checks_groups_and_graph_on_every_rank(): - evidence = _validate_parallel_evidence( - [_hybrid_tp4_ep2_summary(rank) for rank in range(4)], - tensor_parallel_size=4, - expert_parallel_size=2, - last_moe_layer=46, - require_decode_cuda_graph=True, - require_decode_cuda_graph_execution=True, - ) - - assert evidence["hybrid_moe"] is True - assert evidence["local_expert_ranges"] == [ - [0, 32], - [0, 32], - [32, 64], - [32, 64], - ] - assert evidence["expected_groups"]["moe_tensor"] == [[0, 1], [2, 3]] From e253062e32e3d619ed2fb1c2aa713f9643d961d4 Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Wed, 12 Aug 2026 16:13:21 +0800 Subject: [PATCH 27/48] docs: generalize glm setup guidance --- README.md | 1 + requirements/locks/README.md | 16 ++++------------ requirements/locks/canonical-cu129-py310.txt | 2 +- 3 files changed, 6 insertions(+), 13 deletions(-) diff --git a/README.md b/README.md index d7505202..ba7057fb 100644 --- a/README.md +++ b/README.md @@ -69,6 +69,7 @@ Read the method overview and integration rules in | Qwen3 | ✅ | | Qwen3MoE | ✅ | | Qwen3.5 / Qwen3.6 | ✅ | +| Qwen3.5 / Qwen3.6 MoE | ✅ | | GLM-4.7-Flash | ✅ | | Llama 3 / 3.1 | ✅ | | MiniMax M2.7 | ✅ | diff --git a/requirements/locks/README.md b/requirements/locks/README.md index 0e79c61a..540b7136 100644 --- a/requirements/locks/README.md +++ b/requirements/locks/README.md @@ -7,20 +7,12 @@ ranges. Training, benchmark, and test packages are part of the main install. The lock remains the only fully validated package combination; versions inside the wider ranges are candidates until they pass the same validation gates. -Create the environment on shared machines under `/data2`: +Create an isolated environment from the repository root: ```bash -CONDA_PKGS_DIRS=/data2/$USER/cache/conda_pkgs \ - conda create -p /data2/$USER/conda_envs/sparse-vllm-cu129-py310 \ - python=3.10 pip setuptools wheel -y - -PYTHONNOUSERSITE=1 PIP_CACHE_DIR=/data2/$USER/cache/pip \ - /data2/$USER/conda_envs/sparse-vllm-cu129-py310/bin/python -s -m pip \ - install -r requirements/locks/canonical-cu129-py310.txt - -PYTHONNOUSERSITE=1 \ - /data2/$USER/conda_envs/sparse-vllm-cu129-py310/bin/python -s -m pip \ - install --no-deps -e . +uv venv --python 3.10 +uv pip install -r requirements/locks/canonical-cu129-py310.txt +uv pip install --no-deps -e . ``` The lock is the compatibility contract. An unlocked install is a development diff --git a/requirements/locks/canonical-cu129-py310.txt b/requirements/locks/canonical-cu129-py310.txt index 4c2a234e..57e372f6 100644 --- a/requirements/locks/canonical-cu129-py310.txt +++ b/requirements/locks/canonical-cu129-py310.txt @@ -1,5 +1,5 @@ # Canonical main-dependency lock for Python 3.10 and CUDA 12.9. -# Validated on 2026-08-08 with GLM-4.7-Flash TP=2 and CUDA Graph enabled. +# Validated with GLM-4.7-Flash TP=2 and CUDA Graph enabled. # Keep the CUDA-specific local versions; install from the listed indexes. --index-url https://pypi.mirrors.ustc.edu.cn/simple --extra-index-url https://download.pytorch.org/whl/cu129 From cbb70ff595e694e9aa75a40cf3c4f64e8e852098 Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Wed, 12 Aug 2026 16:26:22 +0800 Subject: [PATCH 28/48] test: inject glm cpu router provider --- tests/test_glm4_moe_lite.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/tests/test_glm4_moe_lite.py b/tests/test_glm4_moe_lite.py index bac09c85..bc7b0588 100644 --- a/tests/test_glm4_moe_lite.py +++ b/tests/test_glm4_moe_lite.py @@ -34,6 +34,7 @@ from sparsevllm.models.qwen3 import Qwen3MLP from sparsevllm.operators.mla_attention import MlaAttentionOpSpec from sparsevllm.operators.moe import TritonMoeProvider +from sparsevllm.operators.moe_router import GlmBiasedSigmoidRouterProvider from sparsevllm.platforms import device_runtime @@ -137,6 +138,12 @@ def _construction_context(context: ParallelContext): return_value=TritonMoeProvider(), ) ) + stack.enter_context( + patch( + "sparsevllm.models.glm4_moe_lite.resolve_moe_router_provider", + return_value=GlmBiasedSigmoidRouterProvider(), + ) + ) stack.enter_context(torch.device("cpu")) previous_dtype = torch.get_default_dtype() torch.set_default_dtype(torch.bfloat16) @@ -322,7 +329,11 @@ def test_glm_attention_binds_key_materializer_once_per_manager_context() -> None def test_glm_router_uses_bias_only_for_selection_and_scales_weights() -> None: torch.manual_seed(23) - router = Glm4MoeLiteRouter(_config()) + with patch( + "sparsevllm.models.glm4_moe_lite.resolve_moe_router_provider", + return_value=GlmBiasedSigmoidRouterProvider(), + ): + router = Glm4MoeLiteRouter(_config()) router.weight.data.normal_(mean=0.0, std=0.1) router.e_score_correction_bias.data.zero_() router.e_score_correction_bias.data[:4] = 100.0 From c0103fcf408804cfa075df6f0a38705759812e35 Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Wed, 12 Aug 2026 17:40:31 +0800 Subject: [PATCH 29/48] fix: allow glm graph without eager mode --- src/sparsevllm/method_registry.py | 2 +- tests/test_glm_runtime_compatibility.py | 5 ----- 2 files changed, 1 insertion(+), 6 deletions(-) diff --git a/src/sparsevllm/method_registry.py b/src/sparsevllm/method_registry.py index ecfecb0f..f2447f1e 100644 --- a/src/sparsevllm/method_registry.py +++ b/src/sparsevllm/method_registry.py @@ -138,7 +138,7 @@ class ModelRuntimeCompatibility: prefix_cache_methods=frozenset( {"", "streamingllm", "snapkv", "h2o", "omnikv", "rkv"} ), - requires_eager=True, + requires_eager=False, decode_cuda_graph_methods=frozenset( {"", "streamingllm", "snapkv", "h2o", "omnikv", "rkv"} ), diff --git a/tests/test_glm_runtime_compatibility.py b/tests/test_glm_runtime_compatibility.py index db748c3a..51efadc4 100644 --- a/tests/test_glm_runtime_compatibility.py +++ b/tests/test_glm_runtime_compatibility.py @@ -29,11 +29,6 @@ def test_mla_prefill_workspace_budget_must_be_positive(): ("override", "error_type", "message"), [ ({"vllm_sparse_method": "quest"}, ValueError, "Unsupported glm4_moe_lite"), - ( - {"enforce_eager": False}, - ValueError, - "requires enforce_eager=True", - ), ], ) def test_glm_config_rejects_unsupported_storage_combinations( From dbae852f72d8991f1efecdf0c9ec943698bd7515 Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Wed, 12 Aug 2026 17:54:56 +0800 Subject: [PATCH 30/48] perf: unify all-reduce providers --- .../flashinfer_custom_all_reduce.py | 242 -------------- .../distributed/parallel_context.py | 66 +--- src/sparsevllm/engine/decode_cuda_graph.py | 19 +- src/sparsevllm/engine/model_runner.py | 34 +- src/sparsevllm/layers/embed_head.py | 4 +- src/sparsevllm/models/glm4_moe_lite.py | 157 ++++++++- src/sparsevllm/models/minimax_m2.py | 14 +- src/sparsevllm/operators/all_reduce.py | 301 +++++++++++++++--- tests/test_glm4_moe_lite.py | 5 + tests/test_operator_providers.py | 75 ++++- tests/test_parallel_context.py | 160 ---------- tests/test_tp_rpc.py | 2 - 12 files changed, 531 insertions(+), 548 deletions(-) delete mode 100644 src/sparsevllm/distributed/flashinfer_custom_all_reduce.py diff --git a/src/sparsevllm/distributed/flashinfer_custom_all_reduce.py b/src/sparsevllm/distributed/flashinfer_custom_all_reduce.py deleted file mode 100644 index d52cf5bf..00000000 --- a/src/sparsevllm/distributed/flashinfer_custom_all_reduce.py +++ /dev/null @@ -1,242 +0,0 @@ -from __future__ import annotations - -import ctypes -from contextlib import contextmanager -from typing import TYPE_CHECKING, Iterator - -import torch -import torch.distributed as dist - -from sparsevllm.utils.log import logger - -if TYPE_CHECKING: - from sparsevllm.distributed.parallel_context import ParallelGroup - - -class FlashInferCustomAllReduce: - """Small-tensor TP all-reduce using FlashInfer's vLLM kernel.""" - - _REQUIRED_COMM_APIS = ( - "CudaRTLibrary", - "create_shared_buffer", - "vllm_all_reduce", - "vllm_dispose", - "vllm_get_graph_buffer_ipc_meta", - "vllm_init_custom_ar", - "vllm_meta_size", - "vllm_register_buffer", - "vllm_register_graph_buffers", - ) - - @classmethod - def unsupported_reason(cls, group: ParallelGroup) -> str | None: - """Return why this group cannot use the provider, or ``None``.""" - - if group.size != 2 or group.process_group is None: - return "requires a two-rank process group" - if dist.get_backend(group.process_group) != dist.Backend.NCCL: - return "requires an NCCL process group" - if not torch.cuda.is_available(): - return "requires CUDA" - - device_count = int(torch.cuda.device_count()) - if any(rank < 0 or rank >= device_count for rank in group.ranks): - return ( - "requires all group ranks to map to visible local CUDA devices, " - f"got ranks={group.ranks}, device_count={device_count}" - ) - current_device = int(torch.cuda.current_device()) - capability = tuple( - int(value) for value in torch.cuda.get_device_capability(current_device) - ) - if capability != (9, 0): - return f"requires CUDA SM90, got {capability}" - expected_device = int(group.ranks[group.rank]) - if current_device != expected_device: - return ( - "requires world-rank-to-device mapping for CUDA IPC, " - f"got current_device={current_device}, expected={expected_device}" - ) - for peer_rank, peer_device in enumerate(group.ranks): - if peer_rank == group.rank: - continue - if not torch.cuda.can_device_access_peer(current_device, peer_device): - return ( - "requires CUDA peer access between all ranks, " - f"missing {current_device}->{peer_device} for ranks={group.ranks}" - ) - - try: - from flashinfer import comm - except (ImportError, OSError, RuntimeError) as exc: - return f"FlashInfer communication APIs are unavailable: {exc}" - missing = [name for name in cls._REQUIRED_COMM_APIS if not hasattr(comm, name)] - if missing: - return "FlashInfer communication APIs are missing: " + ", ".join(missing) - return None - - def __init__( - self, - group: ParallelGroup, - *, - max_size_bytes: int = 8 * 1024 * 1024, - num_ctas: int = 32, - ) -> None: - if max_size_bytes <= 0 or num_ctas <= 0: - raise ValueError( - "FlashInfer custom all-reduce sizes must be positive, got " - f"max_size_bytes={max_size_bytes}, num_ctas={num_ctas}." - ) - unsupported_reason = self.unsupported_reason(group) - if unsupported_reason is not None: - raise RuntimeError( - "FlashInfer custom all-reduce is unsupported: " - f"{unsupported_reason}." - ) - - from flashinfer.comm import ( - CudaRTLibrary, - create_shared_buffer, - vllm_all_reduce, - vllm_dispose, - vllm_get_graph_buffer_ipc_meta, - vllm_init_custom_ar, - vllm_meta_size, - vllm_register_buffer, - vllm_register_graph_buffers, - ) - - self.group = group - self.max_size_bytes = int(max_size_bytes) - self.num_ctas = int(num_ctas) - self._cudart = CudaRTLibrary() - self._cuda_ipc_close = self._cudart.lib.cudaIpcCloseMemHandle - self._cuda_ipc_close.restype = ctypes.c_int - self._cuda_ipc_close.argtypes = [ctypes.c_void_p] - self._all_reduce = vllm_all_reduce - self._dispose = vllm_dispose - self._get_graph_buffer_ipc_meta = vllm_get_graph_buffer_ipc_meta - self._register_graph_buffers = vllm_register_graph_buffers - self._capturing = False - self._closed = False - self._supported_sizes: dict[int, int] = {} - self._fallback_sizes: dict[int, int] = {} - self._captured_sizes: dict[int, int] = {} - - self._meta_ptrs = create_shared_buffer( - vllm_meta_size() + self.max_size_bytes, - group.process_group, - ) - self._buffer_ptrs = create_shared_buffer( - self.max_size_bytes, - group.process_group, - ) - self._rank_data = torch.empty( - 8 * 1024 * 1024, - dtype=torch.uint8, - device="cuda", - ) - # TP2 is valid without assuming NVLink. Peer access was checked above. - self._handle = vllm_init_custom_ar( - self._meta_ptrs, - self._rank_data, - group.rank, - False, - ) - vllm_register_buffer(self._handle, self._buffer_ptrs) - - def supports(self, tensor: torch.Tensor) -> bool: - size_bytes = int(tensor.numel() * tensor.element_size()) - return bool( - tensor.is_cuda - and tensor.dtype in {torch.bfloat16, torch.float16} - and tensor.is_contiguous() - and size_bytes > 0 - and size_bytes % 16 == 0 - and size_bytes <= self.max_size_bytes - ) - - def all_reduce(self, tensor: torch.Tensor) -> torch.Tensor | None: - size_bytes = int(tensor.numel() * tensor.element_size()) - if not self.supports(tensor): - self._fallback_sizes[size_bytes] = self._fallback_sizes.get(size_bytes, 0) + 1 - return None - self._supported_sizes[size_bytes] = self._supported_sizes.get(size_bytes, 0) + 1 - output = torch.empty_like(tensor) - registered = self._capturing and torch.cuda.is_current_stream_capturing() - if registered: - self._captured_sizes[size_bytes] = self._captured_sizes.get(size_bytes, 0) + 1 - self._all_reduce( - self._handle, - tensor, - output, - 0 if registered else self._buffer_ptrs[self.group.rank], - 0 if registered else self.max_size_bytes, - self.num_ctas, - ) - return output - - def _register_captured_buffers(self) -> None: - handles, offsets = self._get_graph_buffer_ipc_meta(self._handle) - metadata = [None] * self.group.size - dist.all_gather_object( - metadata, - [handles, offsets], - group=self.group.process_group, - ) - self._register_graph_buffers( - self._handle, - [item[0] for item in metadata], - [item[1] for item in metadata], - ) - logger.info( - "FlashInfer custom TP all-reduce capture evidence: rank={}, " - "captured_sizes={}, supported_sizes={}, fallback_sizes={}, " - "graph_buffers={}", - self.group.rank, - dict(sorted(self._captured_sizes.items())), - dict(sorted(self._supported_sizes.items())), - dict(sorted(self._fallback_sizes.items())), - len(offsets), - ) - - @contextmanager - def capture(self) -> Iterator[None]: - if self._capturing: - raise RuntimeError("FlashInfer custom all-reduce capture is already active.") - self._capturing = True - completed = False - try: - yield - completed = True - finally: - self._capturing = False - if completed: - self._register_captured_buffers() - - def close(self) -> None: - if self._closed: - return - self._closed = True - self._dispose(self._handle) - self._close_shared_buffer(self._buffer_ptrs) - self._close_shared_buffer(self._meta_ptrs) - - def _close_shared_buffer(self, pointers: list[int]) -> None: - device_ids = [torch.cuda.current_device()] - dist.barrier(group=self.group.process_group, device_ids=device_ids) - for peer_rank, pointer in enumerate(pointers): - if peer_rank == self.group.rank: - continue - result = int(self._cuda_ipc_close(ctypes.c_void_p(pointer))) - if result != 0: - raise RuntimeError( - "cudaIpcCloseMemHandle failed: " - f"{self._cudart.cudaGetErrorString(result)}" - ) - dist.barrier(group=self.group.process_group, device_ids=device_ids) - self._cudart.cudaFree(ctypes.c_void_p(pointers[self.group.rank])) - dist.barrier(group=self.group.process_group, device_ids=device_ids) - - -__all__ = ["FlashInferCustomAllReduce"] diff --git a/src/sparsevllm/distributed/parallel_context.py b/src/sparsevllm/distributed/parallel_context.py index 262c4227..608a407c 100644 --- a/src/sparsevllm/distributed/parallel_context.py +++ b/src/sparsevllm/distributed/parallel_context.py @@ -1,26 +1,11 @@ from __future__ import annotations -from contextlib import contextmanager -from dataclasses import dataclass, replace -from typing import Iterator, Protocol +from dataclasses import dataclass import torch import torch.distributed as dist from sparsevllm.distributed.topology import ParallelTopology, parallel_group_ranks -from sparsevllm.utils.log import logger - - -class GraphAllReduceProvider(Protocol): - group: ParallelGroup - - def all_reduce(self, tensor: torch.Tensor) -> torch.Tensor | None: ... - - def capture(self): ... - - def close(self) -> None: ... - - @dataclass(frozen=True) class ParallelGroup: process_group: dist.ProcessGroup | None @@ -46,7 +31,6 @@ class ParallelContext: expert: ParallelGroup data: ParallelGroup moe_tensor: ParallelGroup | None = None - all_reduce_provider: GraphAllReduceProvider | None = None @property def world_rank(self) -> int: @@ -137,15 +121,6 @@ def _all_reduce_out_of_place( group: ParallelGroup, op: dist.ReduceOp, ) -> torch.Tensor: - provider = self.all_reduce_provider - if ( - provider is not None - and provider.group.ranks == group.ranks - and op == dist.ReduceOp.SUM - ): - output = provider.all_reduce(tensor) - if output is not None: - return output return self._all_reduce(tensor, group, op) def tp_all_reduce_out_of_place( @@ -162,19 +137,6 @@ def world_all_reduce_out_of_place( ) -> torch.Tensor: return self._all_reduce_out_of_place(tensor, self.world, op) - @contextmanager - def all_reduce_capture(self) -> Iterator[None]: - provider = self.all_reduce_provider - if provider is None: - yield - return - with provider.capture(): - yield - - def close_all_reduce_provider(self) -> None: - if self.all_reduce_provider is not None: - self.all_reduce_provider.close() - def ep_all_reduce( self, tensor: torch.Tensor, @@ -257,7 +219,6 @@ def _local_group( def init_parallel_context( *, topology: ParallelTopology, - enable_flashinfer_custom_all_reduce: bool = False, ) -> ParallelContext: global _PARALLEL_CONTEXT if _PARALLEL_CONTEXT is not None: @@ -303,31 +264,6 @@ def init_parallel_context( ), ) _PARALLEL_CONTEXT = context - if enable_flashinfer_custom_all_reduce: - from sparsevllm.distributed.flashinfer_custom_all_reduce import ( - FlashInferCustomAllReduce, - ) - - group = _PARALLEL_CONTEXT.tensor - local_reason = FlashInferCustomAllReduce.unsupported_reason(group) - rank_reasons = [None] * group.size - dist.all_gather_object(rank_reasons, local_reason, group=group.process_group) - unsupported = [ - (rank, reason) - for rank, reason in zip(group.ranks, rank_reasons) - if reason is not None - ] - if unsupported: - logger.info( - "Using NCCL TP all-reduce because FlashInfer custom all-reduce " - "is unsupported on ranks {}.", - unsupported, - ) - else: - _PARALLEL_CONTEXT = replace( - _PARALLEL_CONTEXT, - all_reduce_provider=FlashInferCustomAllReduce(group), - ) return _PARALLEL_CONTEXT diff --git a/src/sparsevllm/engine/decode_cuda_graph.py b/src/sparsevllm/engine/decode_cuda_graph.py index e94a7344..a950bbe9 100644 --- a/src/sparsevllm/engine/decode_cuda_graph.py +++ b/src/sparsevllm/engine/decode_cuda_graph.py @@ -438,16 +438,15 @@ def _capture( graph_input_sparse_state_refs = self._snapshot_sparse_state_refs() graph = torch.cuda.CUDAGraph() try: - with self.parallel_context.all_reduce_capture(): - with torch.cuda.graph(graph, pool=self.graph_pool): - self._reset_graph_input_attn_scores(graph_input_sparse_state_refs) - logits = self.run_model(input_ids, positions, is_prefill=False) - if state.key.capture_sampling: - if logits is None: - raise RuntimeError("decode_cuda_graph capture_sampling requires rank-0 logits.") - token_ids = logits.argmax(dim=-1) - else: - token_ids = None + with torch.cuda.graph(graph, pool=self.graph_pool): + self._reset_graph_input_attn_scores(graph_input_sparse_state_refs) + logits = self.run_model(input_ids, positions, is_prefill=False) + if state.key.capture_sampling: + if logits is None: + raise RuntimeError("decode_cuda_graph capture_sampling requires rank-0 logits.") + token_ids = logits.argmax(dim=-1) + else: + token_ids = None except Exception as exc: raise RuntimeError(f"decode_cuda_graph capture failed: {exc!r}") from exc diff --git a/src/sparsevllm/engine/model_runner.py b/src/sparsevllm/engine/model_runner.py index e77b2829..d9d783dd 100644 --- a/src/sparsevllm/engine/model_runner.py +++ b/src/sparsevllm/engine/model_runner.py @@ -46,9 +46,13 @@ Qwen3MoeForCausalLM = None try: - from sparsevllm.models.glm4_moe_lite import Glm4MoeLiteForCausalLM + from sparsevllm.models.glm4_moe_lite import ( + Glm4MoeLiteForCausalLM, + build_glm4_moe_lite_runtime_config, + ) except ImportError: Glm4MoeLiteForCausalLM = None + build_glm4_moe_lite_runtime_config = None try: from sparsevllm.models.minimax_m2 import ( @@ -154,18 +158,7 @@ def __init__( ) self.parallel_context = init_parallel_context( topology=config.parallel_topology, - enable_flashinfer_custom_all_reduce=( - hf_config.model_type == "glm4_moe_lite" - and bool(config.decode_cuda_graph) - and config.tensor_parallel_size == 2 - and config.data_parallel_size == 1 - and self.platform.is_cuda() - ), ) - if self.parallel_context.all_reduce_provider is not None: - logger.info( - "Enabled FlashInfer custom TP all-reduce for tensors up to 8 MiB." - ) # CUDA allocator peaks are process-global and survive LLMEngine.exit(). # Start a new lifecycle before model construction so KV sizing observes # only this engine's model load and persistent allocations. @@ -212,6 +205,20 @@ def __init__( "decode_cuda_graph": config.decode_cuda_graph, "expect_mtp_weights": not config.tiny_random, } + if self.parallel_context.world_size > 1: + if build_glm4_moe_lite_runtime_config is None: + raise ImportError("GLM runtime operators are unavailable.") + model_kwargs["runtime_config"] = build_glm4_moe_lite_runtime_config( + hf_config, + self.parallel_context, + max_decode_tokens=_resolve_decode_static_batch_capacity( + decode_static_capture_sizes, + max_num_seqs_in_batch=config.max_num_seqs_in_batch, + max_decoding_seqs=config.max_decoding_seqs, + ), + cuda_graph=config.decode_cuda_graph, + device_index=int(self.device.index or 0), + ) elif config.model_spec.runtime_class_name == "MiniMaxM2ForCausalLM": if build_minimax_m2_runtime_config is None: raise ImportError("MiniMax M2 runtime operators are unavailable.") @@ -227,7 +234,7 @@ def __init__( cuda_graph=config.decode_cuda_graph, device_index=int(self.device.index or 0), ) - + self.model = _create_model(hf_config, config.model_spec, **model_kwargs) if config.tiny_random: from sparsevllm.debug.tiny_random import initialize_sparse_model @@ -387,7 +394,6 @@ def exit(self): ) if self.rank == 0: self.shm.unlink() - self.parallel_context.close_all_reduce_provider() reset_parallel_context() dist.destroy_process_group() diff --git a/src/sparsevllm/layers/embed_head.py b/src/sparsevllm/layers/embed_head.py index e8645acd..7531fb0a 100644 --- a/src/sparsevllm/layers/embed_head.py +++ b/src/sparsevllm/layers/embed_head.py @@ -12,6 +12,7 @@ def __init__( self, num_embeddings: int, embedding_dim: int, + reduce_results: bool = True, ): super().__init__() self.parallel_context = get_parallel_context() @@ -19,6 +20,7 @@ def __init__( self.tp_size = self.parallel_context.tp_size assert num_embeddings % self.tp_size == 0 self.num_embeddings = num_embeddings + self.reduce_results = bool(reduce_results) self.num_embeddings_per_partition = self.num_embeddings // self.tp_size self.vocab_start_idx = self.num_embeddings_per_partition * self.tp_rank self.vocab_end_idx = self.vocab_start_idx + self.num_embeddings_per_partition @@ -56,7 +58,7 @@ def forward(self, x: torch.Tensor): y = F.embedding(x, self.weight) if self.tp_size > 1: y = mask.unsqueeze(1) * y - return self.parallel_context.tp_all_reduce_out_of_place(y) + return self.parallel_context.tp_all_reduce_out_of_place(y) if self.reduce_results else y class ParallelLMHead(VocabParallelEmbedding): diff --git a/src/sparsevllm/models/glm4_moe_lite.py b/src/sparsevllm/models/glm4_moe_lite.py index b8516598..18d5e615 100644 --- a/src/sparsevllm/models/glm4_moe_lite.py +++ b/src/sparsevllm/models/glm4_moe_lite.py @@ -2,13 +2,14 @@ import os import re +from dataclasses import dataclass import torch import torch.nn.functional as F from torch import nn from transformers import Glm4MoeLiteConfig -from sparsevllm.distributed import get_parallel_context +from sparsevllm.distributed import ParallelContext, ParallelGroup, get_parallel_context from sparsevllm.models.layout import resolve_attention_qk_head_dim from sparsevllm.engine.cache_manager import ( AttentionKeyComputeView, @@ -30,6 +31,11 @@ from sparsevllm.layers.rotary_embedding import RotaryEmbedding, get_rope from sparsevllm.models.qwen3 import Qwen3MLP from sparsevllm.operators.mla_attention import MlaAttentionOpSpec +from sparsevllm.operators.all_reduce import ( + AllReduceOpSpec, + PreparedAllReduceOp, + prepare_all_reduce_op, +) from sparsevllm.operators.activation import resolve_silu_and_mul_provider from sparsevllm.operators.moe import ( MoeOpSpec, @@ -62,6 +68,82 @@ _MTP_PREFIX = f"model.layers.{_MTP_LAYER_INDEX}." +@dataclass +class Glm4MoeLiteRuntimeConfig: + attention_decode_all_reduce: PreparedAllReduceOp + moe_decode_all_reduce: PreparedAllReduceOp + _closed: bool = False + + def close(self) -> None: + if self._closed: + return + ops = (self.attention_decode_all_reduce, self.moe_decode_all_reduce) + for op in {id(op): op for op in ops}.values(): + op.close() + self._closed = True + + +def _prepare_decode_all_reduce( + group: ParallelGroup, + *, + max_rows: int, + hidden_size: int, + dtype: torch.dtype, + cuda_graph: bool, + device_index: int, +) -> PreparedAllReduceOp: + return prepare_all_reduce_op( + AllReduceOpSpec( + world_size=group.size, + ranks=group.ranks, + max_rows=max_rows, + hidden_size=hidden_size, + dtype=dtype, + cuda_graph=cuda_graph, + backend=( + "none" + if group.process_group is None + else str(torch.distributed.get_backend(group.process_group)) + ), + ), + group=group.process_group, + rank=group.rank, + device_index=device_index, + ) + + +def build_glm4_moe_lite_runtime_config( + config: Glm4MoeLiteConfig, + parallel_context: ParallelContext, + *, + max_decode_tokens: int, + cuda_graph: bool, + device_index: int, +) -> Glm4MoeLiteRuntimeConfig: + max_decode_tokens = int(max_decode_tokens) + moe_op = _prepare_decode_all_reduce( + parallel_context.world, + max_rows=2 * max_decode_tokens, + hidden_size=int(config.hidden_size), + dtype=model_activation_dtype(config), + cuda_graph=cuda_graph, + device_index=device_index, + ) + attention_op = ( + moe_op + if parallel_context.attention.ranks == parallel_context.world.ranks + else _prepare_decode_all_reduce( + parallel_context.attention, + max_rows=max_decode_tokens, + hidden_size=int(config.hidden_size), + dtype=model_activation_dtype(config), + cuda_graph=cuda_graph, + device_index=device_index, + ) + ) + return Glm4MoeLiteRuntimeConfig(attention_op, moe_op) + + def build_glm4_moe_lite_mla_attention( config: Glm4MoeLiteConfig, *, @@ -105,10 +187,12 @@ def __init__( mla_attention: MLAAttention, *, projection_chunk_size: int, + runtime_config: Glm4MoeLiteRuntimeConfig | None = None, ) -> None: super().__init__() self.mla_attention = mla_attention self.parallel_context = get_parallel_context() + self.runtime_config = runtime_config self.num_heads = int(config.num_attention_heads) self.local_heads = int(mla_attention.spec.local_q_heads) self.q_lora_rank = int(config.q_lora_rank) @@ -164,6 +248,7 @@ def __init__( self.num_heads * self.v_head_dim, int(config.hidden_size), bias=bool(config.attention_bias), + reduce_results=runtime_config is None, ) self._attention_key_materializer_binding: tuple[int, int] | None = None @@ -384,7 +469,12 @@ def forward( latent, k_rope, ) - return self._project_output(value_output, hidden_states) + output = self._project_output(value_output, hidden_states) + if self.runtime_config is None: + return output + if get_context().is_prefill: + return self.parallel_context.attention_tp_all_reduce(output) + return self.runtime_config.attention_decode_all_reduce.run(output) class Glm4MoeLiteRouter(nn.Module): @@ -555,9 +645,11 @@ def __init__( *, mlp_chunk_size: int, decode_cuda_graph: bool, + runtime_config: Glm4MoeLiteRuntimeConfig | None = None, ) -> None: super().__init__() self.parallel_context = get_parallel_context() + self.runtime_config = runtime_config self.mlp_chunk_size = int(mlp_chunk_size) if self.mlp_chunk_size <= 0: raise ValueError( @@ -639,9 +731,9 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: ], dim=0, ) - return self.parallel_context.world_all_reduce_out_of_place( - local_output - ) + if self.runtime_config is not None: + return self.runtime_config.moe_decode_all_reduce.run(local_output) + return self.parallel_context.world_all_reduce_out_of_place(local_output) if not debug_enabled: if int(hidden_states.shape[0]) <= self.mlp_chunk_size: routed = self._routed_chunk(hidden_states) @@ -705,13 +797,22 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: # Hybrid TP+EP makes both branches partial over the same # outer world. Sum them locally so one collective completes # routed experts and the TP-sharded shared expert together. - output = self.parallel_context.world_all_reduce_out_of_place( - routed + shared_local + local_output = routed + shared_local + output = ( + self.runtime_config.moe_decode_all_reduce.run(local_output) + if self.runtime_config is not None and not context.is_prefill + else self.parallel_context.world_all_reduce_out_of_place( + local_output + ) ) else: # Retain the pure-EP semantic path for direct module use: the # shared expert is replicated rather than TP-sharded. - self.parallel_context.world_all_reduce(routed) + routed = ( + self.runtime_config.moe_decode_all_reduce.run(routed) + if self.runtime_config is not None and not context.is_prefill + else self.parallel_context.world_all_reduce(routed) + ) output = routed + shared_local else: shared_local = self._shared_chunk(hidden_states) @@ -726,8 +827,13 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: # Decode tensors are small. Pack both partials into one # collective, then add the independently reduced rows. This # preserves the original BF16 reduction/addition order. - partials = self.parallel_context.world_all_reduce_out_of_place( - torch.stack((routed, shared_local), dim=0) + partials = torch.stack((routed, shared_local), dim=0) + partials = ( + self.runtime_config.moe_decode_all_reduce.run(partials) + if self.runtime_config is not None + else self.parallel_context.world_all_reduce_out_of_place( + partials + ) ) output = partials[0] + partials[1] if debug_enabled: @@ -747,13 +853,16 @@ def __init__( *, mlp_chunk_size: int, decode_cuda_graph: bool, + runtime_config: Glm4MoeLiteRuntimeConfig | None = None, ) -> None: super().__init__() self.parallel_context = get_parallel_context() + self.runtime_config = runtime_config self.self_attn = Glm4MoeLiteAttention( config, mla_attention, projection_chunk_size=mlp_chunk_size, + runtime_config=runtime_config, ) layer_types = list(config.mlp_layer_types) if len(layer_types) != int(config.num_hidden_layers): @@ -769,6 +878,7 @@ def __init__( hidden_act=str(config.hidden_act), mlp_chunk_size=int(mlp_chunk_size), quantization=None, + reduce_results=runtime_config is None, activation_provider=resolve_silu_and_mul_provider( activation_dtype=model_activation_dtype(config), ), @@ -778,6 +888,7 @@ def __init__( config, mlp_chunk_size=mlp_chunk_size, decode_cuda_graph=decode_cuda_graph, + runtime_config=runtime_config, ) else: raise ValueError( @@ -810,6 +921,12 @@ def forward( residual, ) hidden_states = self.mlp(hidden_states) + if self.runtime_config is not None and isinstance(self.mlp, Qwen3MLP): + hidden_states = ( + self.parallel_context.attention_tp_all_reduce(hidden_states) + if get_context().is_prefill + else self.runtime_config.attention_decode_all_reduce.run(hidden_states) + ) return hidden_states, residual @@ -821,13 +938,17 @@ def __init__( *, mlp_chunk_size: int, decode_cuda_graph: bool, + runtime_config: Glm4MoeLiteRuntimeConfig | None = None, ) -> None: super().__init__() self.config = config self.mla_attention = mla_attention + self.parallel_context = get_parallel_context() + self.runtime_config = runtime_config self.embed_tokens = VocabParallelEmbedding( config.vocab_size, config.hidden_size, + reduce_results=runtime_config is None, ) rope_parameters = getattr(config, "rope_parameters", None) or {} self.rotary_emb = get_rope( @@ -847,6 +968,7 @@ def __init__( mla_attention, mlp_chunk_size=mlp_chunk_size, decode_cuda_graph=decode_cuda_graph, + runtime_config=runtime_config, ) for layer_idx in range(int(config.num_hidden_layers)) ] @@ -859,9 +981,15 @@ def forward( input_ids: torch.Tensor, positions: torch.Tensor, ) -> torch.Tensor: + context = get_context() hidden_states = self.embed_tokens(input_ids) + if self.runtime_config is not None: + hidden_states = ( + self.parallel_context.attention_tp_all_reduce(hidden_states) + if context.is_prefill + else self.runtime_config.attention_decode_all_reduce.run(hidden_states) + ) residual = None - context = get_context() debug_layers_env = os.getenv("SPARSEVLLM_DEBUG_HIDDEN_LAYERS") debug_layers = None if debug_layers_env: @@ -957,15 +1085,18 @@ def __init__( mlp_chunk_size: int, decode_cuda_graph: bool, expect_mtp_weights: bool, + runtime_config: Glm4MoeLiteRuntimeConfig | None = None, ) -> None: super().__init__() self.config = config self.parallel_context = get_parallel_context() + self.runtime_config = runtime_config self.model = Glm4MoeLiteModel( config, mla_attention, mlp_chunk_size=mlp_chunk_size, decode_cuda_graph=decode_cuda_graph, + runtime_config=runtime_config, ) self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size) if config.tie_word_embeddings: @@ -974,6 +1105,10 @@ def __init__( self._intentionally_skipped_expert_weights: set[str] = set() self._intentionally_skipped_mtp_weights: set[str] = set() + def close_runtime_operators(self) -> None: + if self.runtime_config is not None: + self.runtime_config.close() + def _sparse_block(self, layer_idx: int) -> Glm4MoeLiteSparseMoeBlock: if not 0 <= int(layer_idx) < len(self.model.layers): raise ValueError( diff --git a/src/sparsevllm/models/minimax_m2.py b/src/sparsevllm/models/minimax_m2.py index a7e2383f..33f90294 100644 --- a/src/sparsevllm/models/minimax_m2.py +++ b/src/sparsevllm/models/minimax_m2.py @@ -88,7 +88,7 @@ def close(self) -> None: def _prepare_decode_all_reduce( group: ParallelGroup, *, - max_tokens: int, + max_rows: int, hidden_size: int, dtype: torch.dtype, cuda_graph: bool, @@ -97,10 +97,16 @@ def _prepare_decode_all_reduce( return prepare_all_reduce_op( AllReduceOpSpec( world_size=group.size, - max_tokens=max_tokens, + ranks=group.ranks, + max_rows=max_rows, hidden_size=hidden_size, dtype=dtype, cuda_graph=cuda_graph, + backend=( + "none" + if group.process_group is None + else str(torch.distributed.get_backend(group.process_group)) + ), ), group=group.process_group, rank=group.rank, @@ -155,7 +161,7 @@ def build_minimax_m2_runtime_config( ) moe_decode_all_reduce = _prepare_decode_all_reduce( parallel_context.world, - max_tokens=int(max_decode_tokens), + max_rows=int(max_decode_tokens), hidden_size=int(config.hidden_size), dtype=activation_dtype, cuda_graph=bool(cuda_graph), @@ -166,7 +172,7 @@ def build_minimax_m2_runtime_config( else: attention_decode_all_reduce = _prepare_decode_all_reduce( parallel_context.attention, - max_tokens=int(max_decode_tokens), + max_rows=int(max_decode_tokens), hidden_size=int(config.hidden_size), dtype=activation_dtype, cuda_graph=bool(cuda_graph), diff --git a/src/sparsevllm/operators/all_reduce.py b/src/sparsevllm/operators/all_reduce.py index e359ce3f..c2314216 100644 --- a/src/sparsevllm/operators/all_reduce.py +++ b/src/sparsevllm/operators/all_reduce.py @@ -1,5 +1,6 @@ from __future__ import annotations +import ctypes import re from dataclasses import dataclass from importlib.metadata import PackageNotFoundError, version @@ -21,14 +22,23 @@ @dataclass(frozen=True) class AllReduceOpSpec: world_size: int - max_tokens: int + ranks: tuple[int, ...] + max_rows: int hidden_size: int dtype: torch.dtype cuda_graph: bool + backend: str def __post_init__(self) -> None: - if self.world_size <= 0 or self.max_tokens <= 0 or self.hidden_size <= 0: + if self.world_size <= 0 or self.max_rows <= 0 or self.hidden_size <= 0: raise ValueError("All-reduce dimensions must be positive.") + if ( + len(self.ranks) != self.world_size + or len(set(self.ranks)) != self.world_size + ): + raise ValueError( + f"All-reduce ranks must contain world_size unique ranks, got {self.ranks}." + ) class AllReduceProvider: @@ -64,6 +74,40 @@ def run( ) +@dataclass(frozen=True) +class _FlashInferTrtllmProfile: + max_rows: int + launch_with_pdl: bool = False + completion_row_threshold: int | None = None + provider_output_buffer: bool = True + + +_FLASHINFER_TRTLLM_PROFILES = { + ("NVIDIA H100 80GB HBM3", 2, 2048): _FlashInferTrtllmProfile( + max_rows=256, + launch_with_pdl=True, + completion_row_threshold=16, + provider_output_buffer=False, + ), + ("NVIDIA H100 80GB HBM3", 4, 3072): _FlashInferTrtllmProfile(max_rows=32), +} + + +def _flashinfer_dependency_reason() -> str | None: + if find_spec("flashinfer") is None: + return "flashinfer is not installed" + try: + installed = version("flashinfer-python") + except PackageNotFoundError: + return "flashinfer-python package metadata is unavailable" + numeric = tuple(int(part) for part in re.findall(r"\d+", installed)[:3]) + return ( + f"requires flashinfer-python >= 0.6.15, got {installed}" + if numeric < (0, 6, 15) + else None + ) + + @ALL_REDUCE_REGISTRY.register class FlashInferTrtllmAllReduceProvider(AllReduceProvider): name = "flashinfer_trtllm_sm90" @@ -87,35 +131,28 @@ def supports(cls, spec: AllReduceOpSpec, caps: DeviceCaps) -> SupportResult: ) if not caps.supports_graph_capture or not spec.cuda_graph: return SupportResult.no("requires CUDA Graph execution") - if spec.world_size != 4: - return SupportResult.no( - f"requires profiled world_size=4, got {spec.world_size}" - ) - if spec.max_tokens > 32: - return SupportResult.no( - f"requires max_tokens <= 32, got {spec.max_tokens}" - ) - if spec.hidden_size != 3072 or spec.dtype != torch.bfloat16: + if spec.backend != "nccl": + return SupportResult.no(f"requires NCCL, got {spec.backend}") + profile = _FLASHINFER_TRTLLM_PROFILES.get( + (caps.device_name, spec.world_size, spec.hidden_size) + ) + if profile is None or spec.dtype != torch.bfloat16: return SupportResult.no( - "requires profiled BF16 hidden_size=3072, got " - f"{spec.dtype} hidden_size={spec.hidden_size}" + "requires a profiled BF16 topology/shape, got " + f"world_size={spec.world_size} hidden_size={spec.hidden_size} " + f"dtype={spec.dtype}" ) - if find_spec("flashinfer") is None: - return SupportResult.no("flashinfer is not installed") - try: - installed = version("flashinfer-python") - except PackageNotFoundError: - return SupportResult.no("flashinfer-python package metadata is unavailable") - numeric = tuple(int(part) for part in re.findall(r"\d+", installed)[:3]) - if numeric < (0, 6, 15): + if spec.max_rows > profile.max_rows: return SupportResult.no( - f"requires flashinfer-python >= 0.6.15, got {installed}" + f"requires max_rows <= {profile.max_rows}, got {spec.max_rows}" ) - return SupportResult.yes() + reason = _flashinfer_dependency_reason() + return SupportResult.no(reason) if reason else SupportResult.yes() def __init__(self) -> None: self.workspace = None self._output_buffer: torch.Tensor | None = None + self._profile: _FlashInferTrtllmProfile | None = None def prepare(self, spec, *, group, rank, device_index=None) -> None: from flashinfer.comm import create_allreduce_fusion_workspace @@ -124,6 +161,8 @@ def prepare(self, spec, *, group, rank, device_index=None) -> None: raise RuntimeError("FlashInfer all-reduce provider is already prepared.") if group is None: raise RuntimeError("FlashInfer all-reduce requires a distributed process group.") + if dist.get_backend(group) != dist.Backend.NCCL: + raise RuntimeError("FlashInfer all-reduce requires an NCCL process group.") current_device = torch.cuda.current_device() if device_index is None: device_index = current_device @@ -132,37 +171,47 @@ def prepare(self, spec, *, group, rank, device_index=None) -> None: "FlashInfer all-reduce must be prepared on the selected CUDA device: " f"selected={device_index} current={current_device}." ) + caps = platforms.current_platform.get_device_caps(current_device) + profile = _FLASHINFER_TRTLLM_PROFILES[ + (caps.device_name, spec.world_size, spec.hidden_size) + ] workspace = create_allreduce_fusion_workspace( backend="trtllm", world_size=spec.world_size, rank=rank, - max_token_num=spec.max_tokens, + max_token_num=profile.max_rows, hidden_dim=spec.hidden_size, dtype=spec.dtype, group=group, ) try: - output_buffer = torch.empty( - (spec.max_tokens, spec.hidden_size), - dtype=spec.dtype, - device=torch.device("cuda", int(device_index)), + output_buffer = ( + torch.empty( + (spec.max_rows, spec.hidden_size), + dtype=spec.dtype, + device=torch.device("cuda", int(device_index)), + ) + if profile.provider_output_buffer + else None ) except Exception: workspace.destroy() raise self.workspace = workspace self._output_buffer = output_buffer + self._profile = profile def close(self) -> None: workspace = self.workspace self.workspace = None self._output_buffer = None + self._profile = None if workspace is not None: workspace.destroy() def run(self, spec, tensor, *, group) -> torch.Tensor: del group - if self.workspace is None or self._output_buffer is None: + if self.workspace is None or self._profile is None: raise RuntimeError("FlashInfer all-reduce provider was not prepared.") if not tensor.is_cuda: raise ValueError( @@ -171,15 +220,190 @@ def run(self, spec, tensor, *, group) -> torch.Tensor: from flashinfer.comm import allreduce_fusion from flashinfer.comm.trtllm_ar import AllReduceFusionPattern - output = self._output_buffer[: int(tensor.shape[0])] - allreduce_fusion( - input=tensor, + flattened = tensor.view(-1, spec.hidden_size) + output = ( + None + if self._output_buffer is None + else self._output_buffer[: int(flattened.shape[0])] + ) + result = allreduce_fusion( + input=flattened, workspace=self.workspace, pattern=AllReduceFusionPattern.kAllReduce, + launch_with_pdl=self._profile.launch_with_pdl, + trigger_completion_at_end=( + self._profile.completion_row_threshold is None + or int(flattened.shape[0]) > self._profile.completion_row_threshold + ), output=output, ) + return result.view_as(tensor) + + +@ALL_REDUCE_REGISTRY.register +class FlashInferVllmAllReduceProvider(AllReduceProvider): + name = "flashinfer_vllm_sm90" + priority = 100 + max_rows = 256 + num_ctas = 32 + + @classmethod + def supports(cls, spec: AllReduceOpSpec, caps: DeviceCaps) -> SupportResult: + if caps.platform != PlatformEnum.CUDA or caps.compute_capability != (9, 0): + return SupportResult.no( + f"requires CUDA SM90, got {caps.platform.name} {caps.compute_capability}" + ) + if caps.device_name != "NVIDIA H100 80GB HBM3": + return SupportResult.no( + "requires profiled NVIDIA H100 80GB HBM3 hardware, " + f"got {caps.device_name}" + ) + if spec.cuda_graph: + return SupportResult.no("requires eager execution") + if spec.backend != "nccl": + return SupportResult.no(f"requires NCCL, got {spec.backend}") + if ( + spec.world_size != 2 + or spec.max_rows > cls.max_rows + or spec.hidden_size != 2048 + or spec.dtype != torch.bfloat16 + ): + return SupportResult.no( + "requires profiled TP2 BF16 [..., 2048] with max_rows <= 256, " + f"got world_size={spec.world_size} max_rows={spec.max_rows} " + f"hidden_size={spec.hidden_size} dtype={spec.dtype}" + ) + reason = _flashinfer_dependency_reason() + if reason: + return SupportResult.no(reason) + try: + from flashinfer import comm + except (ImportError, OSError, RuntimeError) as exc: + return SupportResult.no( + f"FlashInfer communication APIs are unavailable: {exc}" + ) + required = ( + "CudaRTLibrary", + "create_shared_buffer", + "vllm_all_reduce", + "vllm_dispose", + "vllm_init_custom_ar", + "vllm_meta_size", + "vllm_register_buffer", + ) + missing = [name for name in required if not hasattr(comm, name)] + return ( + SupportResult.no( + "FlashInfer communication APIs are missing: " + ", ".join(missing) + ) + if missing + else SupportResult.yes() + ) + + def __init__(self) -> None: + self._group: dist.ProcessGroup | None = None + self._rank = -1 + self._max_size_bytes = 0 + self._rank_data: torch.Tensor | None = None + self._meta_ptrs: list[int] = [] + self._buffer_ptrs: list[int] = [] + self._handle = None + self._cudart = None + + def prepare(self, spec, *, group, rank, device_index=None) -> None: + from flashinfer.comm import ( + CudaRTLibrary, + create_shared_buffer, + vllm_init_custom_ar, + vllm_meta_size, + vllm_register_buffer, + ) + + if self._handle is not None: + raise RuntimeError("FlashInfer all-reduce provider is already prepared.") + if group is None: + raise RuntimeError( + "FlashInfer all-reduce requires a distributed process group." + ) + current_device = torch.cuda.current_device() + if device_index is None: + device_index = current_device + if int(device_index) != current_device or current_device != spec.ranks[rank]: + raise RuntimeError( + "FlashInfer vLLM all-reduce requires rank-to-device mapping: " + f"ranks={spec.ranks} rank={rank} selected={device_index} " + f"current={current_device}." + ) + if any( + peer != current_device + and not torch.cuda.can_device_access_peer(current_device, peer) + for peer in spec.ranks + ): + raise RuntimeError("FlashInfer vLLM all-reduce requires CUDA peer access.") + max_size_bytes = spec.max_rows * spec.hidden_size * spec.dtype.itemsize + meta_ptrs = create_shared_buffer(vllm_meta_size() + max_size_bytes, group) + buffer_ptrs = create_shared_buffer(max_size_bytes, group) + rank_data = torch.empty(8 * 1024 * 1024, dtype=torch.uint8, device="cuda") + handle = vllm_init_custom_ar(meta_ptrs, rank_data, rank, False) + vllm_register_buffer(handle, buffer_ptrs) + self._group = group + self._rank = rank + self._max_size_bytes = max_size_bytes + self._rank_data = rank_data + self._meta_ptrs = meta_ptrs + self._buffer_ptrs = buffer_ptrs + self._handle = handle + self._cudart = CudaRTLibrary() + + def run(self, spec, tensor, *, group) -> torch.Tensor: + del spec, group + if self._handle is None: + raise RuntimeError("FlashInfer all-reduce provider was not prepared.") + from flashinfer.comm import vllm_all_reduce + + output = torch.empty_like(tensor) + vllm_all_reduce( + self._handle, + tensor, + output, + self._buffer_ptrs[self._rank], + self._max_size_bytes, + self.num_ctas, + ) return output + @staticmethod + def _close_shared_buffer(pointers, group, rank, cudart) -> None: + dist.barrier(group=group, device_ids=[torch.cuda.current_device()]) + close = cudart.lib.cudaIpcCloseMemHandle + close.restype = ctypes.c_int + close.argtypes = [ctypes.c_void_p] + for peer_rank, pointer in enumerate(pointers): + if peer_rank != rank: + result = int(close(ctypes.c_void_p(pointer))) + if result != 0: + raise RuntimeError( + f"cudaIpcCloseMemHandle failed: {cudart.cudaGetErrorString(result)}" + ) + dist.barrier(group=group, device_ids=[torch.cuda.current_device()]) + cudart.cudaFree(ctypes.c_void_p(pointers[rank])) + dist.barrier(group=group, device_ids=[torch.cuda.current_device()]) + + def close(self) -> None: + if self._handle is None: + return + from flashinfer.comm import vllm_dispose + + vllm_dispose(self._handle) + self._handle = None + self._close_shared_buffer( + self._buffer_ptrs, self._group, self._rank, self._cudart + ) + self._close_shared_buffer( + self._meta_ptrs, self._group, self._rank, self._cudart + ) + self._rank_data = None + @ALL_REDUCE_REGISTRY.register class TorchDistributedAllReduceProvider(AllReduceProvider): @@ -206,15 +430,16 @@ def _validate_tensor_contract(spec: AllReduceOpSpec, tensor: torch.Tensor) -> No raise TypeError( f"All-reduce expected dtype={spec.dtype}, got {tensor.dtype}." ) - if tensor.ndim != 2: + if tensor.ndim < 2: raise ValueError( - f"All-reduce expects [tokens, hidden], got shape={tuple(tensor.shape)}." + f"All-reduce expects [..., hidden], got shape={tuple(tensor.shape)}." ) - token_count, hidden_size = (int(value) for value in tensor.shape) - if not 0 < token_count <= spec.max_tokens: + row_count = tensor.numel() // int(tensor.shape[-1]) + hidden_size = int(tensor.shape[-1]) + if not 0 < row_count <= spec.max_rows: raise ValueError( - "All-reduce token count is outside the prepared range: " - f"tokens={token_count} max_tokens={spec.max_tokens}." + "All-reduce row count is outside the prepared range: " + f"rows={row_count} max_rows={spec.max_rows}." ) if hidden_size != spec.hidden_size: raise ValueError( diff --git a/tests/test_glm4_moe_lite.py b/tests/test_glm4_moe_lite.py index bc7b0588..64998aec 100644 --- a/tests/test_glm4_moe_lite.py +++ b/tests/test_glm4_moe_lite.py @@ -426,6 +426,7 @@ def test_glm_sparse_moe_reduces_pure_ep_over_world(ep_size: int) -> None: block = object.__new__(Glm4MoeLiteSparseMoeBlock) nn.Module.__init__(block) block.parallel_context = context + block.runtime_config = None block.mlp_chunk_size = 8 block.shared_experts = nn.Identity() block._routed_chunk = lambda hidden_states: hidden_states.clone() @@ -458,6 +459,7 @@ def test_glm_sparse_moe_reduces_pure_tp_over_world( block = object.__new__(Glm4MoeLiteSparseMoeBlock) nn.Module.__init__(block) block.parallel_context = context + block.runtime_config = None block.mlp_chunk_size = 8 block.shared_experts = nn.Identity() block._routed_chunk = lambda hidden_states: hidden_states.clone() @@ -492,6 +494,7 @@ def test_glm_sparse_moe_reduces_hybrid_tp_ep_shards_over_outer_world() -> None: block = object.__new__(Glm4MoeLiteSparseMoeBlock) nn.Module.__init__(block) block.parallel_context = context + block.runtime_config = None block.mlp_chunk_size = 8 block.shared_experts = nn.Identity() block._routed_chunk = lambda hidden_states: hidden_states.clone() @@ -531,6 +534,7 @@ def test_glm_moe_debug_contract_populates_model_runner_summaries() -> None: block = object.__new__(Glm4MoeLiteSparseMoeBlock) nn.Module.__init__(block) block.parallel_context = context + block.runtime_config = None block.mlp_chunk_size = 8 class _Gate(nn.Module): @@ -628,6 +632,7 @@ def test_glm_decoder_syncs_replicated_attention_before_post_norm( layer = object.__new__(Glm4MoeLiteDecoderLayer) nn.Module.__init__(layer) layer.parallel_context = context + layer.runtime_config = None class _InputNorm(nn.Module): def forward(self, hidden_states, residual): diff --git a/tests/test_operator_providers.py b/tests/test_operator_providers.py index 9928e7ed..9b06153d 100644 --- a/tests/test_operator_providers.py +++ b/tests/test_operator_providers.py @@ -133,10 +133,12 @@ def _gate_up_spec(**overrides) -> GateUpSwiGLUOpSpec: def _all_reduce_spec(**overrides) -> AllReduceOpSpec: values = { "world_size": 4, - "max_tokens": 8, + "ranks": (0, 1, 2, 3), + "max_rows": 8, "hidden_size": 3072, "dtype": torch.bfloat16, "cuda_graph": True, + "backend": "nccl", } values.update(overrides) return AllReduceOpSpec(**values) @@ -159,6 +161,77 @@ def test_minimax_all_reduce_prefers_flashinfer_for_profiled_h100_shape(): assert resolved.provider.name == "flashinfer_trtllm_sm90" +def test_glm_tp2_all_reduce_prefers_profiled_flashinfer_provider(): + with ( + patch("sparsevllm.operators.all_reduce.find_spec", return_value=object()), + patch("sparsevllm.operators.all_reduce.version", return_value="0.6.15.post1"), + ): + resolved = OpResolver(ALL_REDUCE_REGISTRY).resolve( + _all_reduce_spec( + world_size=2, ranks=(0, 1), max_rows=256, hidden_size=2048 + ), + _cuda_caps( + (9, 0), + runtime_version="12.9", + device_name="NVIDIA H100 80GB HBM3", + ), + ) + + assert resolved.provider.name == "flashinfer_trtllm_sm90" + + +def test_glm_tp2_eager_all_reduce_prefers_vllm_provider(): + required_apis = SimpleNamespace( + CudaRTLibrary=object(), + create_shared_buffer=object(), + vllm_all_reduce=object(), + vllm_dispose=object(), + vllm_init_custom_ar=object(), + vllm_meta_size=object(), + vllm_register_buffer=object(), + ) + with ( + patch("sparsevllm.operators.all_reduce.find_spec", return_value=object()), + patch("sparsevllm.operators.all_reduce.version", return_value="0.6.15.post1"), + patch.dict(sys.modules, {"flashinfer": SimpleNamespace(comm=required_apis)}), + ): + resolved = OpResolver(ALL_REDUCE_REGISTRY).resolve( + _all_reduce_spec( + world_size=2, + ranks=(0, 1), + max_rows=256, + hidden_size=2048, + cuda_graph=False, + ), + _cuda_caps( + (9, 0), + runtime_version="12.9", + device_name="NVIDIA H100 80GB HBM3", + ), + ) + + assert resolved.provider.name == "flashinfer_vllm_sm90" + + +def test_unprofiled_all_reduce_rows_use_torch_provider(): + with ( + patch("sparsevllm.operators.all_reduce.find_spec", return_value=object()), + patch("sparsevllm.operators.all_reduce.version", return_value="0.6.15.post1"), + ): + resolved = OpResolver(ALL_REDUCE_REGISTRY).resolve( + _all_reduce_spec( + world_size=2, ranks=(0, 1), max_rows=257, hidden_size=2048 + ), + _cuda_caps( + (9, 0), + runtime_version="12.9", + device_name="NVIDIA H100 80GB HBM3", + ), + ) + + assert resolved.provider.name == "torch_distributed" + + @pytest.mark.parametrize("tp_size", [1, 2]) def test_h20_gate_up_provider_accepts_profiled_qwen36_shape(tp_size): resolved = OpResolver(GATE_UP_SWIGLU_REGISTRY).resolve( diff --git a/tests/test_parallel_context.py b/tests/test_parallel_context.py index 1ee6e1f8..5193df70 100644 --- a/tests/test_parallel_context.py +++ b/tests/test_parallel_context.py @@ -16,9 +16,6 @@ parallel_ranks_from_world_rank, world_rank_from_parallel_ranks, ) -from sparsevllm.distributed.flashinfer_custom_all_reduce import ( - FlashInferCustomAllReduce, -) from sparsevllm.distributed.parallel_context import ( get_parallel_context, init_parallel_context, @@ -163,88 +160,6 @@ def test_hybrid_moe_parallel_context_uses_explicit_groups(): reset_parallel_context() -@pytest.mark.parametrize( - ("local_reason", "rank_reasons", "expects_provider"), - [ - (None, [None, None], True), - ("peer access is unavailable", ["peer access is unavailable"] * 2, False), - (None, [None, "peer access is unavailable"], False), - ], -) -def test_hybrid_tp2_ep2_resolves_custom_all_reduce_at_init( - local_reason, - rank_reasons, - expects_provider, -): - reset_parallel_context() - provider = Mock() - provider_cls = Mock(return_value=provider) - provider_cls.unsupported_reason.return_value = local_reason - - def all_gather_object(output, _local_reason, *, group): - assert group is dist.group.WORLD - output[:] = rank_reasons - - with ( - patch.object(dist, "is_initialized", return_value=True), - patch.object(dist, "get_world_size", return_value=2), - patch.object(dist, "get_rank", return_value=0), - patch.object(dist, "all_gather_object", side_effect=all_gather_object), - patch( - "sparsevllm.distributed.flashinfer_custom_all_reduce." - "FlashInferCustomAllReduce", - provider_cls, - ), - ): - context = init_parallel_context( - topology=ParallelTopology(2, 2, 1, ParallelMode.OUTER_TP_MOE), - enable_flashinfer_custom_all_reduce=True, - ) - - selected_group = provider_cls.unsupported_reason.call_args.args[0] - assert selected_group.ranks == (0, 1) - assert context.attention.ranks == (0, 1) - assert context.expert.ranks == (0, 1) - assert context.moe_tensor.ranks == (0,) - if expects_provider: - provider_cls.assert_called_once_with(selected_group) - assert context.all_reduce_provider is provider - else: - provider_cls.assert_not_called() - assert context.all_reduce_provider is None - reset_parallel_context() - - -def test_custom_all_reduce_support_rejects_missing_peer_access(): - group = ParallelGroup(object(), (0, 1), 0, 2) - with ( - patch.object(dist, "get_backend", return_value=dist.Backend.NCCL), - patch.object(torch.cuda, "is_available", return_value=True), - patch.object(torch.cuda, "device_count", return_value=2), - patch.object(torch.cuda, "current_device", return_value=0), - patch.object(torch.cuda, "get_device_capability", return_value=(9, 0)), - patch.object(torch.cuda, "can_device_access_peer", return_value=False), - ): - reason = FlashInferCustomAllReduce.unsupported_reason(group) - - assert reason is not None - assert "peer access" in reason - - -def test_custom_all_reduce_support_rejects_non_sm90_rank(): - group = ParallelGroup(object(), (0, 1), 0, 2) - with ( - patch.object(dist, "get_backend", return_value=dist.Backend.NCCL), - patch.object(torch.cuda, "is_available", return_value=True), - patch.object(torch.cuda, "device_count", return_value=2), - patch.object(torch.cuda, "current_device", return_value=0), - patch.object(torch.cuda, "get_device_capability", return_value=(8, 9)), - ): - reason = FlashInferCustomAllReduce.unsupported_reason(group) - - assert reason == "requires CUDA SM90, got (8, 9)" - - def test_parallel_context_lifecycle_and_local_groups(): reset_parallel_context() fake_groups = [] @@ -318,81 +233,6 @@ def test_ep_broadcast_rejects_invalid_source_rank(): context.ep_broadcast(torch.tensor([1.0]), src_ep_rank=4) -def test_small_out_of_place_all_reduce_uses_bound_provider(): - ranks = (0, 1) - group = ParallelGroup(object(), ranks, 0, 2) - reduced = torch.tensor([3.0]) - provider = SimpleNamespace( - group=group, - all_reduce=Mock(return_value=reduced), - ) - context = ParallelContext( - world=group, - tensor=group, - expert=ParallelGroup(None, (0,), 0, 1), - data=ParallelGroup(None, (0,), 0, 1), - all_reduce_provider=provider, - ) - tensor = torch.tensor([1.0]) - - with patch.object(dist, "all_reduce") as nccl_all_reduce: - output = context.tp_all_reduce_out_of_place(tensor) - - assert output is reduced - provider.all_reduce.assert_called_once_with(tensor) - nccl_all_reduce.assert_not_called() - - -def test_small_out_of_place_all_reduce_falls_back_to_nccl(): - ranks = (0, 1) - group = ParallelGroup(object(), ranks, 0, 2) - provider = SimpleNamespace(group=group, all_reduce=Mock(return_value=None)) - context = ParallelContext( - world=group, - tensor=group, - expert=ParallelGroup(None, (0,), 0, 1), - data=ParallelGroup(None, (0,), 0, 1), - all_reduce_provider=provider, - ) - tensor = torch.tensor([1.0]) - - with patch.object(dist, "all_reduce") as nccl_all_reduce: - output = context.world_all_reduce_out_of_place(tensor) - - assert output is tensor - nccl_all_reduce.assert_called_once_with( - tensor, - op=dist.ReduceOp.SUM, - group=group.process_group, - ) - - -def test_all_reduce_capture_and_close_delegate_to_provider(): - group = ParallelGroup(None, (0,), 0, 1) - capture = Mock() - capture.return_value.__enter__ = Mock(return_value=None) - capture.return_value.__exit__ = Mock(return_value=False) - provider = SimpleNamespace( - group=group, - capture=capture, - close=Mock(), - ) - context = ParallelContext( - world=group, - tensor=group, - expert=group, - data=group, - all_reduce_provider=provider, - ) - - with context.all_reduce_capture(): - pass - context.close_all_reduce_provider() - - capture.assert_called_once_with() - provider.close.assert_called_once_with() - - @pytest.mark.parametrize("op", [dist.ReduceOp.SUM, dist.ReduceOp.MAX]) def test_parallel_context_collectives_are_always_in_place_torch_operations(op): world_group = object() diff --git a/tests/test_tp_rpc.py b/tests/test_tp_rpc.py index 65cb8f83..b5725d54 100644 --- a/tests/test_tp_rpc.py +++ b/tests/test_tp_rpc.py @@ -454,7 +454,6 @@ def test_model_runner_exit_drains_graphs_before_barrier(): ) runner.parallel_context = SimpleNamespace( world_barrier=lambda **_: calls.append("barrier"), - close_all_reduce_provider=lambda: calls.append("close_all_reduce"), ) with ( @@ -478,7 +477,6 @@ def test_model_runner_exit_drains_graphs_before_barrier(): "close_shm", "barrier", "unlink_shm", - "close_all_reduce", "reset", "destroy", ] From c30d35f318ca7c433460f1f3412185653bcf9ce9 Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Wed, 12 Aug 2026 17:58:52 +0800 Subject: [PATCH 31/48] refactor: delegate model runtime binding --- src/sparsevllm/engine/model_runner.py | 91 ++++++-------------------- src/sparsevllm/models/glm4_moe_lite.py | 36 ++++++++++ src/sparsevllm/models/minimax_m2.py | 20 ++++++ tests/test_glm4_moe_lite.py | 55 ++++++++++++++++ tests/test_minimax_m2_model.py | 28 ++++++++ tests/test_tp_rpc.py | 26 ++++++++ 6 files changed, 185 insertions(+), 71 deletions(-) diff --git a/src/sparsevllm/engine/model_runner.py b/src/sparsevllm/engine/model_runner.py index d9d783dd..9293c9fc 100644 --- a/src/sparsevllm/engine/model_runner.py +++ b/src/sparsevllm/engine/model_runner.py @@ -46,22 +46,14 @@ Qwen3MoeForCausalLM = None try: - from sparsevllm.models.glm4_moe_lite import ( - Glm4MoeLiteForCausalLM, - build_glm4_moe_lite_runtime_config, - ) + from sparsevllm.models.glm4_moe_lite import Glm4MoeLiteForCausalLM except ImportError: Glm4MoeLiteForCausalLM = None - build_glm4_moe_lite_runtime_config = None try: - from sparsevllm.models.minimax_m2 import ( - MiniMaxM2ForCausalLM, - build_minimax_m2_runtime_config, - ) + from sparsevllm.models.minimax_m2 import MiniMaxM2ForCausalLM except ImportError: MiniMaxM2ForCausalLM = None - build_minimax_m2_runtime_config = None try: from sparsevllm.models.qwen3_5 import Qwen35ForCausalLM @@ -74,12 +66,16 @@ Qwen35MoeForCausalLM = None -def _create_model(hf_config, model_spec: ModelSpec, **kwargs): +def _create_model(hf_config, model_spec: ModelSpec, **runtime_kwargs): class_name = model_spec.runtime_class_name model_class = globals().get(class_name) if model_class is None: raise ImportError(f"{class_name} is unavailable for {model_spec.name}.") - return model_class(hf_config, **kwargs) + builder = getattr(model_class, "build_runtime_kwargs", None) + return model_class( + hf_config, + **(builder(hf_config, **runtime_kwargs) if callable(builder) else {}), + ) TP_SHM_NAME_PREFIX = "sparsevllm_" @@ -168,12 +164,6 @@ def __init__( torch.set_default_dtype(hf_config.torch_dtype) torch.set_default_device(self.device) setattr(hf_config, "mlp_chunk_size", config.mlp_chunk_size) - if hf_config.model_type == "minimax_m2": - setattr( - hf_config, - "prefill_kv_view_layer_invariant", - config.vllm_sparse_method == "", - ) setattr( hf_config, "decode_cuda_graph", @@ -183,59 +173,18 @@ def __init__( config.decode_cuda_graph_capture_sizes, config.max_decoding_seqs, ) - model_kwargs = {} - if config.model_spec.runtime_class_name == "Glm4MoeLiteForCausalLM": - from sparsevllm.models.glm4_moe_lite import ( - build_glm4_moe_lite_mla_attention, - ) - - model_kwargs = { - "mla_attention": build_glm4_moe_lite_mla_attention( - hf_config, - device=self.device, - max_batch_size=max( - config.max_num_seqs_in_batch, - config.max_decoding_seqs, - ), - prefill_workspace_bytes=config.mla_prefill_workspace_bytes, - decode_cuda_graph=config.decode_cuda_graph, - projection_chunk_size=config.mlp_chunk_size, - ), - "mlp_chunk_size": config.mlp_chunk_size, - "decode_cuda_graph": config.decode_cuda_graph, - "expect_mtp_weights": not config.tiny_random, - } - if self.parallel_context.world_size > 1: - if build_glm4_moe_lite_runtime_config is None: - raise ImportError("GLM runtime operators are unavailable.") - model_kwargs["runtime_config"] = build_glm4_moe_lite_runtime_config( - hf_config, - self.parallel_context, - max_decode_tokens=_resolve_decode_static_batch_capacity( - decode_static_capture_sizes, - max_num_seqs_in_batch=config.max_num_seqs_in_batch, - max_decoding_seqs=config.max_decoding_seqs, - ), - cuda_graph=config.decode_cuda_graph, - device_index=int(self.device.index or 0), - ) - elif config.model_spec.runtime_class_name == "MiniMaxM2ForCausalLM": - if build_minimax_m2_runtime_config is None: - raise ImportError("MiniMax M2 runtime operators are unavailable.") - model_kwargs["runtime_config"] = build_minimax_m2_runtime_config( - hf_config, - self.parallel_context, - layer_invariant_page_table=config.vllm_sparse_method == "", - max_decode_tokens=_resolve_decode_static_batch_capacity( - decode_static_capture_sizes, - max_num_seqs_in_batch=config.max_num_seqs_in_batch, - max_decoding_seqs=config.max_decoding_seqs, - ), - cuda_graph=config.decode_cuda_graph, - device_index=int(self.device.index or 0), - ) - - self.model = _create_model(hf_config, config.model_spec, **model_kwargs) + self.model = _create_model( + hf_config, + config.model_spec, + engine_config=config, + parallel_context=self.parallel_context, + device=self.device, + max_decode_tokens=_resolve_decode_static_batch_capacity( + decode_static_capture_sizes, + max_num_seqs_in_batch=config.max_num_seqs_in_batch, + max_decoding_seqs=config.max_decoding_seqs, + ), + ) if config.tiny_random: from sparsevllm.debug.tiny_random import initialize_sparse_model diff --git a/src/sparsevllm/models/glm4_moe_lite.py b/src/sparsevllm/models/glm4_moe_lite.py index 18d5e615..318fd870 100644 --- a/src/sparsevllm/models/glm4_moe_lite.py +++ b/src/sparsevllm/models/glm4_moe_lite.py @@ -1077,6 +1077,42 @@ class Glm4MoeLiteForCausalLM(nn.Module): "up_proj": ("gate_up_proj", 1), } + @staticmethod + def build_runtime_kwargs( + config: Glm4MoeLiteConfig, + *, + engine_config, + parallel_context: ParallelContext, + device: torch.device, + max_decode_tokens: int, + ) -> dict: + decode_cuda_graph = bool(engine_config.decode_cuda_graph) + kwargs = { + "mla_attention": build_glm4_moe_lite_mla_attention( + config, + device=device, + max_batch_size=max( + engine_config.max_num_seqs_in_batch, + engine_config.max_decoding_seqs, + ), + prefill_workspace_bytes=engine_config.mla_prefill_workspace_bytes, + decode_cuda_graph=decode_cuda_graph, + projection_chunk_size=engine_config.mlp_chunk_size, + ), + "mlp_chunk_size": engine_config.mlp_chunk_size, + "decode_cuda_graph": decode_cuda_graph, + "expect_mtp_weights": not engine_config.tiny_random, + } + if parallel_context.world_size > 1: + kwargs["runtime_config"] = build_glm4_moe_lite_runtime_config( + config, + parallel_context, + max_decode_tokens=max_decode_tokens, + cuda_graph=decode_cuda_graph, + device_index=int(device.index or 0), + ) + return kwargs + def __init__( self, config: Glm4MoeLiteConfig, diff --git a/src/sparsevllm/models/minimax_m2.py b/src/sparsevllm/models/minimax_m2.py index 33f90294..1e47b087 100644 --- a/src/sparsevllm/models/minimax_m2.py +++ b/src/sparsevllm/models/minimax_m2.py @@ -619,6 +619,26 @@ class MiniMaxM2ForCausalLM(nn.Module): "v_proj": ("qkv_proj", "v"), } + @staticmethod + def build_runtime_kwargs( + config, + *, + engine_config, + parallel_context: ParallelContext, + device: torch.device, + max_decode_tokens: int, + ) -> dict: + return { + "runtime_config": build_minimax_m2_runtime_config( + config, + parallel_context, + layer_invariant_page_table=engine_config.vllm_sparse_method == "", + max_decode_tokens=max_decode_tokens, + cuda_graph=engine_config.decode_cuda_graph, + device_index=int(device.index or 0), + ) + } + def __init__( self, config, diff --git a/tests/test_glm4_moe_lite.py b/tests/test_glm4_moe_lite.py index 64998aec..a1aea9b5 100644 --- a/tests/test_glm4_moe_lite.py +++ b/tests/test_glm4_moe_lite.py @@ -178,6 +178,61 @@ def test_glm_topology_reuses_one_mla_object_and_qwen_dense_mlp() -> None: assert model.model.layers[1].mlp.experts.op_spec.routing_method == "biased_sigmoid" +def test_glm_runtime_kwargs_bind_model_owned_operators() -> None: + config = _config() + context = _tp_context(tp_size=2) + runtime = SimpleNamespace( + decode_cuda_graph=True, + max_num_seqs_in_batch=4, + max_decoding_seqs=8, + mla_prefill_workspace_bytes=1024, + mlp_chunk_size=16, + tiny_random=False, + ) + mla = object() + all_reduce = object() + with ( + patch( + "sparsevllm.models.glm4_moe_lite.build_glm4_moe_lite_mla_attention", + return_value=mla, + ) as build_mla, + patch( + "sparsevllm.models.glm4_moe_lite.build_glm4_moe_lite_runtime_config", + return_value=all_reduce, + ) as build_all_reduce, + ): + kwargs = Glm4MoeLiteForCausalLM.build_runtime_kwargs( + config, + engine_config=runtime, + parallel_context=context, + device=torch.device("cuda", 1), + max_decode_tokens=8, + ) + + assert kwargs == { + "mla_attention": mla, + "mlp_chunk_size": 16, + "decode_cuda_graph": True, + "expect_mtp_weights": True, + "runtime_config": all_reduce, + } + build_mla.assert_called_once_with( + config, + device=torch.device("cuda", 1), + max_batch_size=8, + prefill_workspace_bytes=1024, + decode_cuda_graph=True, + projection_chunk_size=16, + ) + build_all_reduce.assert_called_once_with( + config, + context, + max_decode_tokens=8, + cuda_graph=True, + device_index=1, + ) + + def test_glm_interleaved_rope_matches_transformers() -> None: torch.manual_seed(13) q = torch.randn(1, 3, 5, 64) diff --git a/tests/test_minimax_m2_model.py b/tests/test_minimax_m2_model.py index 620ae39f..ba3da563 100644 --- a/tests/test_minimax_m2_model.py +++ b/tests/test_minimax_m2_model.py @@ -201,6 +201,34 @@ def test_model_construction_shares_explicit_runtime_operators_across_layers(): all_reduce_op.close.assert_called_once_with() +def test_minimax_runtime_kwargs_bind_model_owned_operators(): + config = _config() + context = _tp_context(0, 1) + runtime = SimpleNamespace(vllm_sparse_method="", decode_cuda_graph=True) + bound = object() + with patch( + "sparsevllm.models.minimax_m2.build_minimax_m2_runtime_config", + return_value=bound, + ) as build: + kwargs = MiniMaxM2ForCausalLM.build_runtime_kwargs( + config, + engine_config=runtime, + parallel_context=context, + device=torch.device("cuda", 1), + max_decode_tokens=8, + ) + + assert kwargs == {"runtime_config": bound} + build.assert_called_once_with( + config, + context, + layer_invariant_page_table=True, + max_decode_tokens=8, + cuda_graph=True, + device_index=1, + ) + + def _random_fp8(shape): return torch.randn(shape).clamp(-4.0, 4.0).to(torch.float8_e4m3fn) diff --git a/tests/test_tp_rpc.py b/tests/test_tp_rpc.py index b5725d54..8c686274 100644 --- a/tests/test_tp_rpc.py +++ b/tests/test_tp_rpc.py @@ -17,11 +17,37 @@ TP_RUN_STATUS_SUCCESS, TP_RPC_STATUS_SYNC_METHODS, TP_SHM_NAME_PREFIX, + _create_model, make_tp_shm_name, ) +from sparsevllm.models.spec import ModelSpec from sparsevllm.operators import registry as operator_registry +def test_create_model_delegates_optional_runtime_binding(): + config = object() + context = object() + + class Model: + @classmethod + def build_runtime_kwargs(cls, hf_config, **runtime_kwargs): + assert hf_config is config + assert runtime_kwargs == {"parallel_context": context} + return {"runtime": "bound"} + + def __init__(self, hf_config, *, runtime): + self.args = hf_config, runtime + + with patch.dict(_create_model.__globals__, {"Model": Model}): + model = _create_model( + config, + ModelSpec("test", runtime_class_name="Model"), + parallel_context=context, + ) + + assert model.args == (config, "bound") + + def test_write_shm_waits_until_worker_reads_command(): ctx = get_context("spawn") command_event = ctx.Event() From 89e9468c9785123e393d402a5753bce69fc01807 Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Wed, 12 Aug 2026 18:36:42 +0800 Subject: [PATCH 32/48] refactor: remove obsolete all-reduce wrappers --- .../distributed/parallel_context.py | 22 ------------------- src/sparsevllm/layers/embed_head.py | 2 +- src/sparsevllm/layers/linear.py | 2 +- src/sparsevllm/models/glm4_moe_lite.py | 14 ++++-------- tests/test_parallel_context.py | 6 ++--- tests/test_qwen3_fp8.py | 1 - 6 files changed, 9 insertions(+), 38 deletions(-) diff --git a/src/sparsevllm/distributed/parallel_context.py b/src/sparsevllm/distributed/parallel_context.py index 608a407c..3afc5a44 100644 --- a/src/sparsevllm/distributed/parallel_context.py +++ b/src/sparsevllm/distributed/parallel_context.py @@ -115,28 +115,6 @@ def tp_all_reduce( ) -> torch.Tensor: return self.attention_tp_all_reduce(tensor, op) - def _all_reduce_out_of_place( - self, - tensor: torch.Tensor, - group: ParallelGroup, - op: dist.ReduceOp, - ) -> torch.Tensor: - return self._all_reduce(tensor, group, op) - - def tp_all_reduce_out_of_place( - self, - tensor: torch.Tensor, - op: dist.ReduceOp = dist.ReduceOp.SUM, - ) -> torch.Tensor: - return self._all_reduce_out_of_place(tensor, self.attention, op) - - def world_all_reduce_out_of_place( - self, - tensor: torch.Tensor, - op: dist.ReduceOp = dist.ReduceOp.SUM, - ) -> torch.Tensor: - return self._all_reduce_out_of_place(tensor, self.world, op) - def ep_all_reduce( self, tensor: torch.Tensor, diff --git a/src/sparsevllm/layers/embed_head.py b/src/sparsevllm/layers/embed_head.py index 7531fb0a..5128ca8b 100644 --- a/src/sparsevllm/layers/embed_head.py +++ b/src/sparsevllm/layers/embed_head.py @@ -58,7 +58,7 @@ def forward(self, x: torch.Tensor): y = F.embedding(x, self.weight) if self.tp_size > 1: y = mask.unsqueeze(1) * y - return self.parallel_context.tp_all_reduce_out_of_place(y) if self.reduce_results else y + return self.parallel_context.tp_all_reduce(y) if self.reduce_results else y class ParallelLMHead(VocabParallelEmbedding): diff --git a/src/sparsevllm/layers/linear.py b/src/sparsevllm/layers/linear.py index 1384cc3e..995f0664 100755 --- a/src/sparsevllm/layers/linear.py +++ b/src/sparsevllm/layers/linear.py @@ -492,5 +492,5 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: else: y = F.linear(x, self.weight, bias) if self.reduce_results: - return self.parallel_context.tp_all_reduce_out_of_place(y) + return self.parallel_context.tp_all_reduce(y) return y diff --git a/src/sparsevllm/models/glm4_moe_lite.py b/src/sparsevllm/models/glm4_moe_lite.py index 318fd870..4670db15 100644 --- a/src/sparsevllm/models/glm4_moe_lite.py +++ b/src/sparsevllm/models/glm4_moe_lite.py @@ -733,7 +733,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: ) if self.runtime_config is not None: return self.runtime_config.moe_decode_all_reduce.run(local_output) - return self.parallel_context.world_all_reduce_out_of_place(local_output) + return self.parallel_context.world_all_reduce(local_output) if not debug_enabled: if int(hidden_states.shape[0]) <= self.mlp_chunk_size: routed = self._routed_chunk(hidden_states) @@ -801,9 +801,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: output = ( self.runtime_config.moe_decode_all_reduce.run(local_output) if self.runtime_config is not None and not context.is_prefill - else self.parallel_context.world_all_reduce_out_of_place( - local_output - ) + else self.parallel_context.world_all_reduce(local_output) ) else: # Retain the pure-EP semantic path for direct module use: the @@ -820,9 +818,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: # Both branches are TP partials. Compose them locally so the # full MoE block needs one collective, matching the fused-MoE # communication contract used by the reference runtime. - output = self.parallel_context.world_all_reduce_out_of_place( - routed + shared_local - ) + output = self.parallel_context.world_all_reduce(routed + shared_local) else: # Decode tensors are small. Pack both partials into one # collective, then add the independently reduced rows. This @@ -831,9 +827,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: partials = ( self.runtime_config.moe_decode_all_reduce.run(partials) if self.runtime_config is not None - else self.parallel_context.world_all_reduce_out_of_place( - partials - ) + else self.parallel_context.world_all_reduce(partials) ) output = partials[0] + partials[1] if debug_enabled: diff --git a/tests/test_parallel_context.py b/tests/test_parallel_context.py index 5193df70..eab58a6a 100644 --- a/tests/test_parallel_context.py +++ b/tests/test_parallel_context.py @@ -412,12 +412,12 @@ def test_dense_layers_use_tp_group_in_replicated_ep_topology(): assert embedding.weight.shape == (32, 8) -def test_vocab_parallel_embedding_uses_out_of_place_reduction(): +def test_vocab_parallel_embedding_reduces_results(): reduced = torch.randn(2, 4) context = SimpleNamespace( tp_rank=0, tp_size=2, - tp_all_reduce_out_of_place=Mock(return_value=reduced), + tp_all_reduce=Mock(return_value=reduced), ) with patch( "sparsevllm.layers.embed_head.get_parallel_context", @@ -428,7 +428,7 @@ def test_vocab_parallel_embedding_uses_out_of_place_reduction(): output = embedding(torch.tensor([0, 5])) assert output is reduced - context.tp_all_reduce_out_of_place.assert_called_once() + context.tp_all_reduce.assert_called_once() def test_cache_kv_heads_depend_on_tp_not_ep(): diff --git a/tests/test_qwen3_fp8.py b/tests/test_qwen3_fp8.py index 02fd59e1..174c6ecc 100644 --- a/tests/test_qwen3_fp8.py +++ b/tests/test_qwen3_fp8.py @@ -16,7 +16,6 @@ def _parallel_context(): tp_rank=0, tp_size=1, tp_all_reduce=lambda tensor: tensor, - tp_all_reduce_out_of_place=lambda tensor: tensor, ) From 695a084419e2c19b3a850f187f3c4bf6a6203885 Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Wed, 12 Aug 2026 19:03:53 +0800 Subject: [PATCH 33/48] refactor: remove model-specific cache fallbacks --- src/sparsevllm/configs/model.py | 3 +++ src/sparsevllm/debug/tiny_random.py | 10 +++++++--- .../engine/cache_manager/storage/__init__.py | 9 +-------- tests/test_attention_cache_storage.py | 18 ------------------ 4 files changed, 11 insertions(+), 29 deletions(-) diff --git a/src/sparsevllm/configs/model.py b/src/sparsevllm/configs/model.py index 3e68ca29..6c335f1e 100644 --- a/src/sparsevllm/configs/model.py +++ b/src/sparsevllm/configs/model.py @@ -108,6 +108,9 @@ def load_and_validate_model(config) -> None: config.tiny_random_overrides = apply_tiny_random_overrides( config.hf_config, config.tiny_random_config, + validate_standard_head_shape=( + model_spec.attention_cache_layout == "explicit_kv" + ), ) log_once( "TINY RANDOM MODE is enabled: checkpoint weights will not be read and " diff --git a/src/sparsevllm/debug/tiny_random.py b/src/sparsevllm/debug/tiny_random.py index c07ee770..bbde1853 100644 --- a/src/sparsevllm/debug/tiny_random.py +++ b/src/sparsevllm/debug/tiny_random.py @@ -85,7 +85,12 @@ def load_tiny_random_overrides(path: str) -> dict[str, int]: return overrides -def apply_tiny_random_overrides(hf_config: Any, path: str) -> dict[str, int]: +def apply_tiny_random_overrides( + hf_config: Any, + path: str, + *, + validate_standard_head_shape: bool = True, +) -> dict[str, int]: overrides = load_tiny_random_overrides(path) original_values = { name: int(getattr(hf_config, name)) @@ -130,8 +135,7 @@ def apply_tiny_random_overrides(hf_config: Any, path: str) -> dict[str, int]: num_heads = int(hf_config.num_attention_heads) num_kv_heads = int(hf_config.num_key_value_heads) head_dim = int(getattr(hf_config, "head_dim", hidden_size // num_heads)) - is_mla = str(getattr(hf_config, "model_type", "") or "") == "glm4_moe_lite" - if not is_mla and hidden_size != num_heads * head_dim: + if validate_standard_head_shape and hidden_size != num_heads * head_dim: raise ValueError( "Tiny random config requires hidden_size == num_attention_heads * head_dim, " f"got {hidden_size} != {num_heads} * {head_dim}." diff --git a/src/sparsevllm/engine/cache_manager/storage/__init__.py b/src/sparsevllm/engine/cache_manager/storage/__init__.py index bb1a86ce..2e564758 100644 --- a/src/sparsevllm/engine/cache_manager/storage/__init__.py +++ b/src/sparsevllm/engine/cache_manager/storage/__init__.py @@ -15,14 +15,7 @@ def create_attention_cache_storage( num_kv_heads: int, head_dim: int, ) -> AttentionCacheStorage: - configured_layout = getattr(config, "attention_cache_layout", None) - if configured_layout is None: - model_type = str(getattr(config.hf_config, "model_type", "") or "") - configured_layout = ( - CacheLayout.MLA_LATENT.value - if model_type == "glm4_moe_lite" - else CacheLayout.EXPLICIT_KV.value - ) + configured_layout = config.attention_cache_layout layout = ( configured_layout if isinstance(configured_layout, CacheLayout) diff --git a/tests/test_attention_cache_storage.py b/tests/test_attention_cache_storage.py index ed10a8bf..9333c967 100644 --- a/tests/test_attention_cache_storage.py +++ b/tests/test_attention_cache_storage.py @@ -98,24 +98,6 @@ def test_storage_factory_uses_configured_layout(): MlaLatentStorage, ) - inferred_mla_config = SimpleNamespace( - hf_config=SimpleNamespace( - model_type="glm4_moe_lite", - torch_dtype=torch.bfloat16, - kv_lora_rank=512, - qk_rope_head_dim=64, - ) - ) - assert isinstance( - create_attention_cache_storage( - inferred_mla_config, - num_kv_heads=4, - head_dim=64, - ), - MlaLatentStorage, - ) - - @pytest.mark.parametrize( ("storage", "num_layers", "num_slots", "expected_shape"), [ From fda46e582c93c3107b269774503147a8aeb6cefa Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Wed, 12 Aug 2026 19:11:31 +0800 Subject: [PATCH 34/48] refactor: organize kernel integrations --- .agents/skills/add-sparse-method/SKILL.md | 2 +- .../add-sparse-method/references/file-map.md | 2 +- .../references/operator-integration.md | 5 ++- .../references/tilelang.md | 2 +- .../references/triton.md | 2 +- .../review-operator-organization/SKILL.md | 4 +- docs/en/design/control-map.md | 4 +- docs/zh/design/control-map.md | 4 +- scripts/profiling/bench_prefill_score.py | 2 +- .../kernel_bench/test_context_flash_attn.py | 2 +- scripts/profiling/kernel_bench/test_decode.py | 10 ++--- .../test_gqa_flash_decoding_score.py | 4 +- .../validation/test_omnikv_fused_compare.py | 2 +- skills/add-sparse-method/SKILL.md | 2 +- .../add-sparse-method/references/file-map.md | 2 +- src/sparsevllm/engine/cache_manager/base.py | 2 +- .../engine/cache_manager/deltakv_base.py | 8 ++-- .../cache_manager/deltakv_less_memory.py | 12 ++--- src/sparsevllm/engine/cache_manager/h2o.py | 2 +- src/sparsevllm/engine/cache_manager/rkv.py | 2 +- src/sparsevllm/engine/cache_manager/snapkv.py | 2 +- .../cache_manager/storage/explicit_kv.py | 2 +- .../cache_manager/storage/mla_latent.py | 4 +- src/sparsevllm/engine/sparse_controller.py | 2 +- src/sparsevllm/kernels/__init__.py | 1 + src/sparsevllm/kernels/external/__init__.py | 1 + .../kernels/external/sgl/__init__.py | 1 + .../external/sgl/fa3.py} | 2 +- .../external/sgl/moe.py} | 43 +++++------------- .../external/sgl/support.py} | 1 - src/sparsevllm/kernels/moe.py | 12 +++++ .../tilelang}/__init__.py | 0 .../tilelang}/mla/LICENSE.tilelang | 0 .../tilelang}/mla/README.md | 0 .../tilelang}/mla/__init__.py | 0 .../tilelang}/mla/decode.py | 0 .../tilelang/mla/runtime.py} | 8 ++-- .../triton}/__init__.py | 0 .../triton}/column_parallel_rmsnorm.py | 0 .../triton}/context_flashattention_nopad.py | 0 .../triton}/deltakv_kernels.py | 0 .../triton}/embedding.py | 0 .../triton}/flash_decoding.py | 0 .../triton}/flash_decoding_stage1.py | 0 .../triton}/flash_decoding_stage2.py | 0 .../triton}/fp8_blockwise.py | 0 .../triton}/gate_up_swiglu.py | 0 .../gqa_decode_flashattention_nopad.py | 0 .../triton}/gqa_flash_decoding.py | 0 .../triton}/gqa_flash_decoding_stage1.py | 0 .../triton}/gqa_flash_decoding_stage2.py | 0 .../triton}/minimax_m2_moe.py | 2 +- .../triton}/minimax_m2_router.py | 2 +- .../triton}/mla/LICENSE.lightllm | 0 .../triton}/mla/README.md | 0 .../triton}/mla/__init__.py | 0 .../triton}/mla/copy_latent.py | 0 .../triton}/mla/decode_schedule.py | 0 .../triton}/mla/decode_stage1.py | 0 .../triton}/mla/decode_stage2.py | 0 .../triton}/mla/gather_latent.py | 0 .../{triton_kernel => kernels/triton}/moe.py | 16 ++----- .../triton}/moe_biased_sigmoid.py | 0 .../triton}/moe_config.py | 0 .../triton}/moe_topk.py | 0 .../triton}/omnikv_fused.py | 0 .../triton}/ppl_fp16_flash_decoding.py | 0 .../triton}/ppl_int4kv_copy_kv.py | 0 .../triton}/ppl_int4kv_flash_decoding.py | 0 .../triton}/ppl_int8kv_flash_decoding.py | 0 .../triton}/ppl_quant_copy_kv.py | 0 .../triton}/prefill_score.py | 0 .../triton}/quant.py | 0 .../triton}/qwen3_5/README.md | 0 .../triton}/qwen3_5/__init__.py | 0 .../triton}/qwen3_5/autotuner.py | 0 .../triton}/qwen3_5/causal_conv1d.py | 0 .../triton}/qwen3_5/fla/__init__.py | 0 .../triton}/qwen3_5/fla/ops/__init__.py | 0 .../triton}/qwen3_5/fla/ops/chunk.py | 0 .../triton}/qwen3_5/fla/ops/chunk_delta_h.py | 2 +- .../triton}/qwen3_5/fla/ops/chunk_o.py | 2 +- .../qwen3_5/fla/ops/chunk_scaled_dot_kkt.py | 2 +- .../triton}/qwen3_5/fla/ops/cumsum.py | 2 +- .../qwen3_5/fla/ops/fused_recurrent.py | 0 .../triton}/qwen3_5/fla/ops/index.py | 0 .../triton}/qwen3_5/fla/ops/l2norm.py | 2 +- .../triton}/qwen3_5/fla/ops/op.py | 0 .../triton}/qwen3_5/fla/ops/solve_tril.py | 0 .../triton}/qwen3_5/fla/ops/utils.py | 0 .../triton}/qwen3_5/fla/ops/wy_fast.py | 0 .../triton}/qwen3_5/fused_gdn_gating.py | 2 +- .../triton}/qwen3_5/gated_rmsnorm.py | 2 +- .../triton}/qwen3_5/gated_shared_add.py | 0 .../triton}/qwen3_5/gdn_decode_pack.py | 0 .../triton}/rmsnorm.py | 0 .../triton}/rotary_emb.py | 0 .../triton}/silu_and_mul.py | 0 .../splitfuse_context_flashattention_nopad.py | 0 .../triton}/store_kvcache.py | 0 .../triton}/token_attention_nopad_att1.py | 0 .../triton}/token_attention_nopad_reduceV.py | 0 .../triton}/token_attention_nopad_softmax.py | 0 .../token_attention_softmax_and_reducev.py | 0 src/sparsevllm/layers/attention_backend.py | 14 +++--- src/sparsevllm/layers/layernorm.py | 4 +- src/sparsevllm/layers/mla_attention.py | 2 +- src/sparsevllm/models/glm4_moe_lite.py | 2 +- src/sparsevllm/models/minimax_m2.py | 2 +- src/sparsevllm/models/qwen3_5.py | 10 ++--- src/sparsevllm/models/qwen3_moe.py | 2 +- src/sparsevllm/operators/activation.py | 2 +- src/sparsevllm/operators/fp8_linear.py | 2 +- src/sparsevllm/operators/gate_up_swiglu.py | 2 +- src/sparsevllm/operators/gated_shared_add.py | 2 +- src/sparsevllm/operators/mla_attention.py | 19 ++++---- src/sparsevllm/operators/moe.py | 45 +++++++++++++++---- src/sparsevllm/operators/moe_router.py | 4 +- src/sparsevllm/operators/prefill_attention.py | 2 +- tests/test_deltakv_less_memory_kernel.py | 10 ++--- tests/test_minimax_m2_attention_graph.py | 6 +-- tests/test_minimax_m2_router.py | 4 +- tests/test_mla_attention_operator.py | 2 +- tests/test_mla_kernels.py | 5 ++- tests/test_moe_config.py | 2 +- tests/test_operator_providers.py | 12 ++--- tests/test_prefill_score_kernel.py | 6 +-- tests/test_qwen35_causal_conv1d.py | 2 +- tests/test_qwen35_hd256_decode_routing.py | 10 ++--- tests/test_qwen3_moe.py | 2 +- tests/test_sgl_fa3.py | 2 +- tests/test_sgl_moe.py | 16 +++---- tests/test_tilelang_mla_kernel.py | 3 +- tests/test_tilelang_mla_operator.py | 12 ++--- tests/test_triton_fp8_operators.py | 6 +-- tests/test_triton_moe.py | 8 ++-- 136 files changed, 213 insertions(+), 197 deletions(-) create mode 100644 src/sparsevllm/kernels/__init__.py create mode 100644 src/sparsevllm/kernels/external/__init__.py create mode 100644 src/sparsevllm/kernels/external/sgl/__init__.py rename src/sparsevllm/{operators/sgl_fa3.py => kernels/external/sgl/fa3.py} (99%) rename src/sparsevllm/{operators/sgl_moe.py => kernels/external/sgl/moe.py} (56%) rename src/sparsevllm/{operators/sgl_kernel.py => kernels/external/sgl/support.py} (99%) create mode 100644 src/sparsevllm/kernels/moe.py rename src/sparsevllm/{tilelang_kernel => kernels/tilelang}/__init__.py (100%) rename src/sparsevllm/{tilelang_kernel => kernels/tilelang}/mla/LICENSE.tilelang (100%) rename src/sparsevllm/{tilelang_kernel => kernels/tilelang}/mla/README.md (100%) rename src/sparsevllm/{tilelang_kernel => kernels/tilelang}/mla/__init__.py (100%) rename src/sparsevllm/{tilelang_kernel => kernels/tilelang}/mla/decode.py (100%) rename src/sparsevllm/{operators/tilelang_mla.py => kernels/tilelang/mla/runtime.py} (99%) rename src/sparsevllm/{triton_kernel => kernels/triton}/__init__.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/column_parallel_rmsnorm.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/context_flashattention_nopad.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/deltakv_kernels.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/embedding.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/flash_decoding.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/flash_decoding_stage1.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/flash_decoding_stage2.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/fp8_blockwise.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/gate_up_swiglu.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/gqa_decode_flashattention_nopad.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/gqa_flash_decoding.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/gqa_flash_decoding_stage1.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/gqa_flash_decoding_stage2.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/minimax_m2_moe.py (99%) rename src/sparsevllm/{triton_kernel => kernels/triton}/minimax_m2_router.py (98%) rename src/sparsevllm/{triton_kernel => kernels/triton}/mla/LICENSE.lightllm (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/mla/README.md (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/mla/__init__.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/mla/copy_latent.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/mla/decode_schedule.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/mla/decode_stage1.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/mla/decode_stage2.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/mla/gather_latent.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/moe.py (99%) rename src/sparsevllm/{triton_kernel => kernels/triton}/moe_biased_sigmoid.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/moe_config.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/moe_topk.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/omnikv_fused.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/ppl_fp16_flash_decoding.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/ppl_int4kv_copy_kv.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/ppl_int4kv_flash_decoding.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/ppl_int8kv_flash_decoding.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/ppl_quant_copy_kv.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/prefill_score.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/quant.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/qwen3_5/README.md (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/qwen3_5/__init__.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/qwen3_5/autotuner.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/qwen3_5/causal_conv1d.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/qwen3_5/fla/__init__.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/qwen3_5/fla/ops/__init__.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/qwen3_5/fla/ops/chunk.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/qwen3_5/fla/ops/chunk_delta_h.py (99%) rename src/sparsevllm/{triton_kernel => kernels/triton}/qwen3_5/fla/ops/chunk_o.py (98%) rename src/sparsevllm/{triton_kernel => kernels/triton}/qwen3_5/fla/ops/chunk_scaled_dot_kkt.py (98%) rename src/sparsevllm/{triton_kernel => kernels/triton}/qwen3_5/fla/ops/cumsum.py (99%) rename src/sparsevllm/{triton_kernel => kernels/triton}/qwen3_5/fla/ops/fused_recurrent.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/qwen3_5/fla/ops/index.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/qwen3_5/fla/ops/l2norm.py (98%) rename src/sparsevllm/{triton_kernel => kernels/triton}/qwen3_5/fla/ops/op.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/qwen3_5/fla/ops/solve_tril.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/qwen3_5/fla/ops/utils.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/qwen3_5/fla/ops/wy_fast.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/qwen3_5/fused_gdn_gating.py (97%) rename src/sparsevllm/{triton_kernel => kernels/triton}/qwen3_5/gated_rmsnorm.py (98%) rename src/sparsevllm/{triton_kernel => kernels/triton}/qwen3_5/gated_shared_add.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/qwen3_5/gdn_decode_pack.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/rmsnorm.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/rotary_emb.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/silu_and_mul.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/splitfuse_context_flashattention_nopad.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/store_kvcache.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/token_attention_nopad_att1.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/token_attention_nopad_reduceV.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/token_attention_nopad_softmax.py (100%) rename src/sparsevllm/{triton_kernel => kernels/triton}/token_attention_softmax_and_reducev.py (100%) diff --git a/.agents/skills/add-sparse-method/SKILL.md b/.agents/skills/add-sparse-method/SKILL.md index 5bcd2393..af866e22 100644 --- a/.agents/skills/add-sparse-method/SKILL.md +++ b/.agents/skills/add-sparse-method/SKILL.md @@ -28,7 +28,7 @@ Follow this placement order. 5. Keep `src/sparsevllm/layers/attention.py` method-agnostic. It may call generic hooks, but should not grow method-specific branches unless adding a new reusable hook. 6. Put cross-layer observation, attention-score collection, or scheduler-facing sparse orchestration in `src/sparsevllm/engine/sparse_controller.py`. 7. Use `src/sparsevllm/utils/` only for truly generic helpers shared by multiple methods. Do not place an entire method implementation there. -8. Add custom kernels under `src/sparsevllm/triton_kernel/` or another explicit runtime module, then call them through the method's cache manager or shared decode path. +8. Add custom kernels under `src/sparsevllm/kernels/triton/` or another explicit runtime module, then call them through the method's cache manager or shared decode path. ## Decision Rules diff --git a/.agents/skills/add-sparse-method/references/file-map.md b/.agents/skills/add-sparse-method/references/file-map.md index 3a864b95..f6c34020 100644 --- a/.agents/skills/add-sparse-method/references/file-map.md +++ b/.agents/skills/add-sparse-method/references/file-map.md @@ -48,7 +48,7 @@ Do not bury a full method implementation in `attention.py`. ## Add Kernel Code Only When Needed -Touch `src/sparsevllm/triton_kernel/` or another explicit kernel module when: +Touch `src/sparsevllm/kernels/triton/` or another explicit kernel module when: - the existing decode or prefill kernels are the bottleneck - the method requires a new layout-aware fused operator diff --git a/.agents/skills/optimize-sparsevllm-kernel/references/operator-integration.md b/.agents/skills/optimize-sparsevllm-kernel/references/operator-integration.md index fc997272..61b5f453 100644 --- a/.agents/skills/optimize-sparsevllm-kernel/references/operator-integration.md +++ b/.agents/skills/optimize-sparsevllm-kernel/references/operator-integration.md @@ -16,9 +16,10 @@ model semantic call -> Triton, TileLang, JIT, or external kernel ``` -- Keep repository-owned Triton kernels under `src/sparsevllm/triton_kernel/`. +- Keep repository-owned Triton kernels under `src/sparsevllm/kernels/triton/`. - Keep repository-owned TileLang kernels under - `src/sparsevllm/tilelang_kernel/`. + `src/sparsevllm/kernels/tilelang/`. +- Keep thin third-party adapters under `src/sparsevllm/kernels/external/`. - Keep support checks, dependency availability, static launch selection, workspaces, and provider binding under `src/sparsevllm/operators/`. - Keep models expressed in semantic operations rather than backend names, diff --git a/.agents/skills/optimize-sparsevllm-kernel/references/tilelang.md b/.agents/skills/optimize-sparsevllm-kernel/references/tilelang.md index 7db96924..a0334855 100644 --- a/.agents/skills/optimize-sparsevllm-kernel/references/tilelang.md +++ b/.agents/skills/optimize-sparsevllm-kernel/references/tilelang.md @@ -1,6 +1,6 @@ # TileLang Kernel Guide -Use this guide for kernels under `src/sparsevllm/tilelang_kernel/` and their +Use this guide for kernels under `src/sparsevllm/kernels/tilelang/` and their lazy adapters under `src/sparsevllm/operators/`. ## Establish the Contract diff --git a/.agents/skills/optimize-sparsevllm-kernel/references/triton.md b/.agents/skills/optimize-sparsevllm-kernel/references/triton.md index 8b948177..c8f1f996 100644 --- a/.agents/skills/optimize-sparsevllm-kernel/references/triton.md +++ b/.agents/skills/optimize-sparsevllm-kernel/references/triton.md @@ -1,7 +1,7 @@ # Triton Kernel Guide Use this guide for repository-owned Triton kernels and wrappers under -`src/sparsevllm/triton_kernel/`. +`src/sparsevllm/kernels/triton/`. ## Inspect Before Editing diff --git a/.agents/skills/review-operator-organization/SKILL.md b/.agents/skills/review-operator-organization/SKILL.md index 18ee667b..feca46a3 100644 --- a/.agents/skills/review-operator-organization/SKILL.md +++ b/.agents/skills/review-operator-organization/SKILL.md @@ -25,7 +25,9 @@ Read all changed and directly coupled files in these areas: - `src/sparsevllm/operators/`: specs, providers, registries, and resolvers. - `src/sparsevllm/platforms/`: platform discovery and `DeviceCaps`. -- `src/sparsevllm/triton_kernel/`: repository-owned implementations. +- `src/sparsevllm/kernels/triton/`: repository-owned Triton implementations. +- `src/sparsevllm/kernels/tilelang/`: repository-owned TileLang implementations. +- `src/sparsevllm/kernels/external/`: thin external-kernel adapters. - Model and loader call sites that construct specs or prepare physical weights. - Dependency declarations and installation documentation for external providers. - Resolver, kernel-equivalence, integration, and model tests. diff --git a/docs/en/design/control-map.md b/docs/en/design/control-map.md index 9151fe7e..fc0f2764 100644 --- a/docs/en/design/control-map.md +++ b/docs/en/design/control-map.md @@ -67,7 +67,9 @@ flowchart TD | `src/sparsevllm/engine/cache_manager/*.py` | Physical/logical KV state for each sparse method. | This is the primary place for sparse-method implementation. | | `src/sparsevllm/engine/sparse_controller.py` | Cross-layer attention-score collection, dynamic token selection, post-forward compression triggers. | Keep persistent method metadata in cache managers, not here. | | `src/sparsevllm/layers/attention.py` | Generic KV store + attention kernel dispatch + hook calls. | Add generic hooks if needed; avoid method-specific branches. | -| `src/sparsevllm/triton_kernel/` | Kernel implementations. | Kernel wrappers should fail fast on invalid shape/dtype assumptions. | +| `src/sparsevllm/kernels/triton/` | Repository-owned Triton kernels. | Kernel wrappers should fail fast on invalid shape/dtype assumptions. | +| `src/sparsevllm/kernels/tilelang/` | Repository-owned TileLang kernels and runtime bindings. | Keep compilation and launch details out of operators. | +| `src/sparsevllm/kernels/external/` | Thin adapters for third-party kernel libraries. | Keep optional imports lazy and validate supported API versions. | | `benchmark/model_adapters/sparsevllm.py` | Shared native text-benchmark generation adapter. | Keep it thin; runtime behavior belongs in `src/sparsevllm/`. | | `benchmark/` and `scripts/` | Evaluation, debugging, analysis, throughput scripts. | Preserve raw outputs, parsed outputs, per-sample status, aggregate metrics, and run info separately. | diff --git a/docs/zh/design/control-map.md b/docs/zh/design/control-map.md index 994e2c9f..ac33c32d 100644 --- a/docs/zh/design/control-map.md +++ b/docs/zh/design/control-map.md @@ -57,7 +57,9 @@ flowchart TD | `src/sparsevllm/engine/cache_manager/*.py` | 各稀疏方法的 physical/logical KV state。 | 稀疏方法实现的主要位置。 | | `src/sparsevllm/engine/sparse_controller.py` | 跨 layer attention-score 收集、动态 token selection、post-forward compression trigger。 | Persistent method metadata 应保存在 cache manager，而不是这里。 | | `src/sparsevllm/layers/attention.py` | 通用 KV store、attention kernel dispatch 和 hook 调用。 | 必要时添加 generic hook；避免方法特定 branch。 | -| `src/sparsevllm/triton_kernel/` | Kernel 实现。 | shape/dtype 假设无效时，kernel wrapper 应快速失败。 | +| `src/sparsevllm/kernels/triton/` | 仓库维护的 Triton kernel。 | shape/dtype 假设无效时，kernel wrapper 应快速失败。 | +| `src/sparsevllm/kernels/tilelang/` | 仓库维护的 TileLang kernel 和 runtime binding。 | 编译和 launch 细节不能放入 operators。 | +| `src/sparsevllm/kernels/external/` | 第三方 kernel 库的薄适配。 | 可选依赖保持惰性导入，并校验支持的 API 版本。 | | `benchmark/model_adapters/sparsevllm.py` | 文本 benchmark 共用的原生 generation adapter。 | 保持轻量；runtime 行为属于 `src/sparsevllm/`。 | | `benchmark/` 和 `scripts/` | 评估、调试、分析和吞吐量脚本。 | 分别保存 raw output、parsed output、per-sample status、aggregate metric 和 run info。 | diff --git a/scripts/profiling/bench_prefill_score.py b/scripts/profiling/bench_prefill_score.py index c9add0de..b8faa849 100644 --- a/scripts/profiling/bench_prefill_score.py +++ b/scripts/profiling/bench_prefill_score.py @@ -5,7 +5,7 @@ import torch -from sparsevllm.triton_kernel.prefill_score import prefill_score_fwd +from sparsevllm.kernels.triton.prefill_score import prefill_score_fwd def _make_case( diff --git a/scripts/profiling/kernel_bench/test_context_flash_attn.py b/scripts/profiling/kernel_bench/test_context_flash_attn.py index 7f46f84a..c78ff8f6 100644 --- a/scripts/profiling/kernel_bench/test_context_flash_attn.py +++ b/scripts/profiling/kernel_bench/test_context_flash_attn.py @@ -6,7 +6,7 @@ import fire from typing import List -from sparsevllm.triton_kernel.context_flashattention_nopad import context_attention_fwd +from sparsevllm.kernels.triton.context_flashattention_nopad import context_attention_fwd def get_block_m(): return 64 if "Tesla" in torch.cuda.get_device_name(0) else 128 diff --git a/scripts/profiling/kernel_bench/test_decode.py b/scripts/profiling/kernel_bench/test_decode.py index 5267f45b..a605ce96 100644 --- a/scripts/profiling/kernel_bench/test_decode.py +++ b/scripts/profiling/kernel_bench/test_decode.py @@ -2,14 +2,14 @@ import time import numpy as np from flash_attn import flash_attn_with_kvcache -from sparsevllm.triton_kernel.flash_decoding_stage1 import ( +from sparsevllm.kernels.triton.flash_decoding_stage1 import ( flash_decode_stage1, flash_decode_stage1_with_score ) -from sparsevllm.triton_kernel.flash_decoding_stage2 import flash_decode_stage2 -from sparsevllm.triton_kernel.gqa_decode_flashattention_nopad import gqa_decode_attention_fwd -from sparsevllm.triton_kernel.gqa_flash_decoding_stage1 import flash_decode_stage1 as gqa_flash_decode_stage1 -from sparsevllm.triton_kernel.gqa_flash_decoding_stage2 import flash_decode_stage2 as gqa_flash_decode_stage2 +from sparsevllm.kernels.triton.flash_decoding_stage2 import flash_decode_stage2 +from sparsevllm.kernels.triton.gqa_decode_flashattention_nopad import gqa_decode_attention_fwd +from sparsevllm.kernels.triton.gqa_flash_decoding_stage1 import flash_decode_stage1 as gqa_flash_decode_stage1 +from sparsevllm.kernels.triton.gqa_flash_decoding_stage2 import flash_decode_stage2 as gqa_flash_decode_stage2 def benchmark_gqa_kernels( batch_sizes=[1, 8, 32], diff --git a/scripts/validation/test_gqa_flash_decoding_score.py b/scripts/validation/test_gqa_flash_decoding_score.py index eb54d396..029efd79 100644 --- a/scripts/validation/test_gqa_flash_decoding_score.py +++ b/scripts/validation/test_gqa_flash_decoding_score.py @@ -1,8 +1,8 @@ import torch import numpy as np # 分别导入两个版本的 kernel -from sparsevllm.triton_kernel.flash_decoding_stage1 import flash_decode_stage1_with_score as flash_decode_v1 -from sparsevllm.triton_kernel.gqa_flash_decoding_stage1 import flash_decode_stage1_with_score as flash_decode_gqa +from sparsevllm.kernels.triton.flash_decoding_stage1 import flash_decode_stage1_with_score as flash_decode_v1 +from sparsevllm.kernels.triton.gqa_flash_decoding_stage1 import flash_decode_stage1_with_score as flash_decode_gqa def test_gqa_flash_decoding_score(): torch.manual_seed(42) diff --git a/scripts/validation/test_omnikv_fused_compare.py b/scripts/validation/test_omnikv_fused_compare.py index 3720a0a0..492573cc 100644 --- a/scripts/validation/test_omnikv_fused_compare.py +++ b/scripts/validation/test_omnikv_fused_compare.py @@ -3,7 +3,7 @@ import time import torch -from sparsevllm.triton_kernel.omnikv_fused import build_omnikv_keep_and_slots +from sparsevllm.kernels.triton.omnikv_fused import build_omnikv_keep_and_slots def reference_build(topk_indices, hist_lens, recent_chunk_lens, buffer_req_to_token_slots, req_indices, num_sink): diff --git a/skills/add-sparse-method/SKILL.md b/skills/add-sparse-method/SKILL.md index 7ccf3033..fff1e608 100644 --- a/skills/add-sparse-method/SKILL.md +++ b/skills/add-sparse-method/SKILL.md @@ -42,7 +42,7 @@ Follow this placement order. 6. Put cross-layer observation, attention-score collection, or scheduler-facing sparse orchestration in `src/sparsevllm/engine/sparse_controller.py`. 7. Put hidden-state capture, activation steering, and per-sequence steering state in `src/sparsevllm/engine/activation_controller.py`, with `SparseController` owning the lifecycle and model files calling only a generic hook. 8. Use `src/sparsevllm/utils/` only for truly generic helpers shared by multiple methods. Do not place an entire method implementation there. -9. Add custom kernels under `src/sparsevllm/triton_kernel/` or another explicit runtime module, then call them through the method's cache manager or shared decode path. +9. Add custom kernels under `src/sparsevllm/kernels/triton/` or another explicit runtime module, then call them through the method's cache manager or shared decode path. ## Decision Rules diff --git a/skills/add-sparse-method/references/file-map.md b/skills/add-sparse-method/references/file-map.md index 4917b3ad..95ada08a 100644 --- a/skills/add-sparse-method/references/file-map.md +++ b/skills/add-sparse-method/references/file-map.md @@ -57,7 +57,7 @@ Do not bury a full method implementation in `attention.py`. ## Add Kernel Code Only When Needed -Touch `src/sparsevllm/triton_kernel/` or another explicit kernel module when: +Touch `src/sparsevllm/kernels/triton/` or another explicit kernel module when: - the existing decode or prefill kernels are the bottleneck - the method requires a new layout-aware fused operator diff --git a/src/sparsevllm/engine/cache_manager/base.py b/src/sparsevllm/engine/cache_manager/base.py index 0b85bf0d..65bcab09 100644 --- a/src/sparsevllm/engine/cache_manager/base.py +++ b/src/sparsevllm/engine/cache_manager/base.py @@ -18,7 +18,7 @@ PREFILL_EXECUTION_RAW_OFFLOAD, ) from sparsevllm.method_registry import SUPPORTED_SPARSE_METHODS, normalize_sparse_method -from sparsevllm.triton_kernel.store_kvcache import store_kvcache +from sparsevllm.kernels.triton.store_kvcache import store_kvcache import sparsevllm.platforms as platforms from sparsevllm.models.layout import resolve_attention_qk_head_dim from sparsevllm.utils.log import logger, log_level diff --git a/src/sparsevllm/engine/cache_manager/deltakv_base.py b/src/sparsevllm/engine/cache_manager/deltakv_base.py index 9687453e..a270c715 100644 --- a/src/sparsevllm/engine/cache_manager/deltakv_base.py +++ b/src/sparsevllm/engine/cache_manager/deltakv_base.py @@ -2752,7 +2752,7 @@ def _deltakv_build_view_and_plan_reconstruct_static( buffers = self._ensure_decode_static_plan_buffers(bsz, k_max, max_s, req_indices.device) active_slots, active_pos, local_req, new_context_lens, no_free_temp_slots, recon_pos, recon_latent, recon_out_slot = buffers - from sparsevllm.triton_kernel.deltakv_kernels import deltakv_static_decode_plan + from sparsevllm.kernels.triton.deltakv_kernels import deltakv_static_decode_plan deltakv_static_decode_plan( raw_slots_map=self.sparse_layer_raw_slots_map, @@ -2967,7 +2967,7 @@ def _deltakv_gather_raw_kv( k_cache: torch.Tensor, v_cache: torch.Tensor, ) -> torch.Tensor: - from sparsevllm.triton_kernel.deltakv_kernels import deltakv_gather_raw_kv_grouped_heads + from sparsevllm.kernels.triton.deltakv_kernels import deltakv_gather_raw_kv_grouped_heads hp = int(getattr(self.config, "deltakv_triton_gather_heads_per_program", 4) or 1) hp = max(1, min(hp, int(self.num_kv_heads))) @@ -2993,7 +2993,7 @@ def _deltakv_reconstruct_writeback( v_cache: torch.Tensor, l_idx: int | None = None, ): - from sparsevllm.triton_kernel.deltakv_kernels import deltakv_reconstruct_writeback_grouped_heads + from sparsevllm.kernels.triton.deltakv_kernels import deltakv_reconstruct_writeback_grouped_heads hp = int(getattr(self.config, "deltakv_triton_reconstruct_heads_per_program", 4) or 1) hp = max(1, min(hp, int(self.num_kv_heads))) @@ -3336,7 +3336,7 @@ def _cluster_compress( new_center_rel=new_center_rel, ) - from sparsevllm.triton_kernel.deltakv_kernels import batch_gather_mean, deltakv_l2_topk_blockwise + from sparsevllm.kernels.triton.deltakv_kernels import batch_gather_mean, deltakv_l2_topk_blockwise with profiler.record("deltakv_cluster_metric"): partial_scores, partial_idx = deltakv_l2_topk_blockwise( diff --git a/src/sparsevllm/engine/cache_manager/deltakv_less_memory.py b/src/sparsevllm/engine/cache_manager/deltakv_less_memory.py index 22dd57a4..35b3042d 100644 --- a/src/sparsevllm/engine/cache_manager/deltakv_less_memory.py +++ b/src/sparsevllm/engine/cache_manager/deltakv_less_memory.py @@ -7,13 +7,13 @@ import torch from sparsevllm.engine.sequence import Sequence -from sparsevllm.triton_kernel.quant import ( +from sparsevllm.kernels.triton.quant import ( triton_dequantize_2d_int4_grouped, triton_quantize_and_pack_2d_int4_grouped, triton_quantize_and_pack_along_last_dim, unpack_quantized_to_16bit, ) -from sparsevllm.triton_kernel.deltakv_kernels import deltakv_materialize_sparse_view +from sparsevllm.kernels.triton.deltakv_kernels import deltakv_materialize_sparse_view from sparsevllm.layers.rotary_embedding import apply_rotary_emb from sparsevllm.platforms import device_runtime from sparsevllm.utils.compressor import create_compressor @@ -2814,7 +2814,7 @@ def _deltakv_reconstruct_writeback( v_cache: torch.Tensor, l_idx: int | None = None, ): - from sparsevllm.triton_kernel.deltakv_kernels import deltakv_reconstruct_writeback_grouped_heads + from sparsevllm.kernels.triton.deltakv_kernels import deltakv_reconstruct_writeback_grouped_heads hp = int(getattr(self.config, "deltakv_triton_reconstruct_heads_per_program", 4) or 1) hp = max(1, min(hp, int(self.num_kv_heads))) @@ -2911,7 +2911,7 @@ def _deltakv_build_view_and_plan_reconstruct_static( buffers = self._ensure_decode_static_plan_buffers(bsz, k_max, max_s, req_indices.device) active_slots, active_pos, local_req, new_context_lens, no_free_temp_slots, recon_pos, recon_latent, recon_out_slot = buffers - from sparsevllm.triton_kernel.deltakv_kernels import deltakv_static_decode_plan + from sparsevllm.kernels.triton.deltakv_kernels import deltakv_static_decode_plan deltakv_static_decode_plan( raw_slots_map=self.sparse_layer_raw_slots_map, @@ -3359,7 +3359,7 @@ def _build_full_layer_quantized_view_static( flat_raw_slots = raw_slots.reshape(-1) total_slots = bsz * max_len - from sparsevllm.triton_kernel.deltakv_kernels import full_layer_copy_raw_or_zero + from sparsevllm.kernels.triton.deltakv_kernels import full_layer_copy_raw_or_zero full_layer_copy_raw_or_zero( raw_k=self.full_kv_cache[0, l_idx], @@ -3423,7 +3423,7 @@ def _dequantize_full_layer_kivi_tokens( if ((local_offsets < 0) | (local_offsets >= self._full_layer_kivi_group_size())).any(): raise RuntimeError("Full-layer KIVI token position is outside its packed block.") - from sparsevllm.triton_kernel.deltakv_kernels import full_layer_kivi_dequant_tokens + from sparsevllm.kernels.triton.deltakv_kernels import full_layer_kivi_dequant_tokens if out_k is None or out_v is None: raise RuntimeError("Full-layer KIVI dequantization requires explicit output buffers.") diff --git a/src/sparsevllm/engine/cache_manager/h2o.py b/src/sparsevllm/engine/cache_manager/h2o.py index 85160f60..7089db33 100644 --- a/src/sparsevllm/engine/cache_manager/h2o.py +++ b/src/sparsevllm/engine/cache_manager/h2o.py @@ -7,7 +7,7 @@ import torch.nn.functional as F from sparsevllm.engine.sequence import Sequence -from sparsevllm.triton_kernel.prefill_score import prefill_score_fwd +from sparsevllm.kernels.triton.prefill_score import prefill_score_fwd from sparsevllm.utils.context import get_context from sparsevllm.utils.profiler import profiler diff --git a/src/sparsevllm/engine/cache_manager/rkv.py b/src/sparsevllm/engine/cache_manager/rkv.py index 3746e81e..4ed49a44 100644 --- a/src/sparsevllm/engine/cache_manager/rkv.py +++ b/src/sparsevllm/engine/cache_manager/rkv.py @@ -5,7 +5,7 @@ from sparsevllm.config import Config from sparsevllm.distributed import ParallelContext from sparsevllm.engine.sequence import Sequence -from sparsevllm.triton_kernel.prefill_score import prefill_score_fwd +from sparsevllm.kernels.triton.prefill_score import prefill_score_fwd from .base import PrefillComputeView from .snapkv import SnapKVCacheManager diff --git a/src/sparsevllm/engine/cache_manager/snapkv.py b/src/sparsevllm/engine/cache_manager/snapkv.py index ac4fe079..39bb73cc 100644 --- a/src/sparsevllm/engine/cache_manager/snapkv.py +++ b/src/sparsevllm/engine/cache_manager/snapkv.py @@ -15,7 +15,7 @@ ) from sparsevllm.method_registry import PREFILL_POLICY_LONG_BS1FULL_SHORT_BATCH from sparsevllm.platforms import device_runtime -from sparsevllm.triton_kernel.prefill_score import prefill_score_fwd +from sparsevllm.kernels.triton.prefill_score import prefill_score_fwd from sparsevllm.utils.context import get_context from sparsevllm.utils.log import logger, log_level from sparsevllm.utils.profiler import profiler diff --git a/src/sparsevllm/engine/cache_manager/storage/explicit_kv.py b/src/sparsevllm/engine/cache_manager/storage/explicit_kv.py index 596c1b8e..83874094 100644 --- a/src/sparsevllm/engine/cache_manager/storage/explicit_kv.py +++ b/src/sparsevllm/engine/cache_manager/storage/explicit_kv.py @@ -2,7 +2,7 @@ import torch -from sparsevllm.triton_kernel.store_kvcache import store_kvcache +from sparsevllm.kernels.triton.store_kvcache import store_kvcache from ..base import AttentionCacheWrite, ExplicitKVPayload, ExplicitKVWrite from .base import CacheLayout diff --git a/src/sparsevllm/engine/cache_manager/storage/mla_latent.py b/src/sparsevllm/engine/cache_manager/storage/mla_latent.py index 048efd45..6c378849 100644 --- a/src/sparsevllm/engine/cache_manager/storage/mla_latent.py +++ b/src/sparsevllm/engine/cache_manager/storage/mla_latent.py @@ -2,11 +2,11 @@ import torch -from sparsevllm.triton_kernel.mla.copy_latent import ( +from sparsevllm.kernels.triton.mla.copy_latent import ( copy_latent_to_cache, validate_copy_slot_mapping, ) -from sparsevllm.triton_kernel.mla.decode_stage1 import MLA_LATENT_DIM, MLA_ROPE_DIM +from sparsevllm.kernels.triton.mla.decode_stage1 import MLA_LATENT_DIM, MLA_ROPE_DIM from ..base import AttentionCacheWrite, MlaLatentPayload, MlaLatentWrite from .base import CacheLayout diff --git a/src/sparsevllm/engine/sparse_controller.py b/src/sparsevllm/engine/sparse_controller.py index 4406aa9b..1f9d7214 100644 --- a/src/sparsevllm/engine/sparse_controller.py +++ b/src/sparsevllm/engine/sparse_controller.py @@ -14,7 +14,7 @@ def build_omnikv_keep_and_slots(*args, **kwargs): - from sparsevllm.triton_kernel.omnikv_fused import build_omnikv_keep_and_slots as _build + from sparsevllm.kernels.triton.omnikv_fused import build_omnikv_keep_and_slots as _build return _build(*args, **kwargs) diff --git a/src/sparsevllm/kernels/__init__.py b/src/sparsevllm/kernels/__init__.py new file mode 100644 index 00000000..cdd8b5db --- /dev/null +++ b/src/sparsevllm/kernels/__init__.py @@ -0,0 +1 @@ +"""Repository-owned and external GPU kernel integrations.""" diff --git a/src/sparsevllm/kernels/external/__init__.py b/src/sparsevllm/kernels/external/__init__.py new file mode 100644 index 00000000..a24e2fb5 --- /dev/null +++ b/src/sparsevllm/kernels/external/__init__.py @@ -0,0 +1 @@ +"""Thin adapters for external kernel libraries.""" diff --git a/src/sparsevllm/kernels/external/sgl/__init__.py b/src/sparsevllm/kernels/external/sgl/__init__.py new file mode 100644 index 00000000..8b0ede10 --- /dev/null +++ b/src/sparsevllm/kernels/external/sgl/__init__.py @@ -0,0 +1 @@ +"""SGL kernel adapters.""" diff --git a/src/sparsevllm/operators/sgl_fa3.py b/src/sparsevllm/kernels/external/sgl/fa3.py similarity index 99% rename from src/sparsevllm/operators/sgl_fa3.py rename to src/sparsevllm/kernels/external/sgl/fa3.py index be378924..45cf1f0e 100644 --- a/src/sparsevllm/operators/sgl_fa3.py +++ b/src/sparsevllm/kernels/external/sgl/fa3.py @@ -5,7 +5,7 @@ import torch -from sparsevllm.operators.sgl_kernel import sgl_kernel_support +from sparsevllm.kernels.external.sgl.support import sgl_kernel_support _COMMON_FWD_ARGUMENTS = ( "q", diff --git a/src/sparsevllm/operators/sgl_moe.py b/src/sparsevllm/kernels/external/sgl/moe.py similarity index 56% rename from src/sparsevllm/operators/sgl_moe.py rename to src/sparsevllm/kernels/external/sgl/moe.py index 85b7cee8..b53296af 100644 --- a/src/sparsevllm/operators/sgl_moe.py +++ b/src/sparsevllm/kernels/external/sgl/moe.py @@ -3,8 +3,8 @@ import torch import triton -from sparsevllm.operators.sgl_kernel import sgl_kernel_support -from sparsevllm.triton_kernel.moe import MoeAlignment +from sparsevllm.kernels.external.sgl.support import sgl_kernel_support +from sparsevllm.kernels.moe import MoeAlignment def sgl_moe_alignment_support() -> tuple[bool, str]: @@ -18,37 +18,18 @@ def sgl_moe_align_block_size( *, block_size: int, num_experts: int, - local_expert_start: int, - local_expert_end: int, ) -> MoeAlignment: - """Build a local MoE assignment with the SGL CUDA kernel.""" + """Group local expert assignments with the SGL CUDA kernel.""" - local_expert_start = int(local_expert_start) - local_expert_end = int(local_expert_end) - num_local_experts = local_expert_end - local_expert_start - if not 0 <= local_expert_start < local_expert_end <= int(num_experts): - raise ValueError( - "Invalid local expert range " - f"[{local_expert_start}, {local_expert_end}) for {num_experts} experts." - ) + num_experts = int(num_experts) + if num_experts <= 0: + raise ValueError(f"SGL MoE alignment requires experts, got {num_experts}.") supported, reason = sgl_moe_alignment_support() if not supported: raise RuntimeError(reason) - local_topk_ids = topk_ids - has_remote_experts = num_local_experts != int(num_experts) - if has_remote_experts: - from sparsevllm.triton_kernel.moe import localize_expert_ids - - local_topk_ids = localize_expert_ids( - topk_ids, - local_expert_start=local_expert_start, - local_expert_end=local_expert_end, - remote_expert_id=num_local_experts, - ) - num_assignments = int(local_topk_ids.numel()) - padding_experts = num_local_experts + int(has_remote_experts) + num_assignments = int(topk_ids.numel()) max_num_tokens_padded = triton.cdiv( - num_assignments + padding_experts * (int(block_size) - 1), + num_assignments + num_experts * (int(block_size) - 1), int(block_size), ) * int(block_size) sorted_token_ids = torch.empty( @@ -67,7 +48,7 @@ def sgl_moe_align_block_size( device=topk_ids.device, ) cumsum_buffer = torch.empty( - num_local_experts + 1, + num_experts + 1, dtype=torch.int32, device=topk_ids.device, ) @@ -76,8 +57,8 @@ def sgl_moe_align_block_size( # sgl-kernel 0.3.x iterates to num_experts - 1. The extra empty logical # expert makes the complete [0, num_experts) range participate. moe_align_block_size( - local_topk_ids, - num_local_experts + 1, + topk_ids, + num_experts + 1, int(block_size), sorted_token_ids, expert_ids, @@ -95,6 +76,6 @@ def sgl_moe_align_block_size( __all__ = [ - "sgl_moe_alignment_support", "sgl_moe_align_block_size", + "sgl_moe_alignment_support", ] diff --git a/src/sparsevllm/operators/sgl_kernel.py b/src/sparsevllm/kernels/external/sgl/support.py similarity index 99% rename from src/sparsevllm/operators/sgl_kernel.py rename to src/sparsevllm/kernels/external/sgl/support.py index 30ab6a57..110761f4 100644 --- a/src/sparsevllm/operators/sgl_kernel.py +++ b/src/sparsevllm/kernels/external/sgl/support.py @@ -4,7 +4,6 @@ import importlib.util import re - _MIN_VERSION = (0, 3, 21) _MAX_VERSION = (0, 4, 0) diff --git a/src/sparsevllm/kernels/moe.py b/src/sparsevllm/kernels/moe.py new file mode 100644 index 00000000..f5b63943 --- /dev/null +++ b/src/sparsevllm/kernels/moe.py @@ -0,0 +1,12 @@ +from dataclasses import dataclass + +import torch + + +@dataclass(frozen=True) +class MoeAlignment: + sorted_token_ids: torch.Tensor | None + expert_ids: torch.Tensor + num_tokens_post_padded: torch.Tensor + block_size: int + naive: bool diff --git a/src/sparsevllm/tilelang_kernel/__init__.py b/src/sparsevllm/kernels/tilelang/__init__.py similarity index 100% rename from src/sparsevllm/tilelang_kernel/__init__.py rename to src/sparsevllm/kernels/tilelang/__init__.py diff --git a/src/sparsevllm/tilelang_kernel/mla/LICENSE.tilelang b/src/sparsevllm/kernels/tilelang/mla/LICENSE.tilelang similarity index 100% rename from src/sparsevllm/tilelang_kernel/mla/LICENSE.tilelang rename to src/sparsevllm/kernels/tilelang/mla/LICENSE.tilelang diff --git a/src/sparsevllm/tilelang_kernel/mla/README.md b/src/sparsevllm/kernels/tilelang/mla/README.md similarity index 100% rename from src/sparsevllm/tilelang_kernel/mla/README.md rename to src/sparsevllm/kernels/tilelang/mla/README.md diff --git a/src/sparsevllm/tilelang_kernel/mla/__init__.py b/src/sparsevllm/kernels/tilelang/mla/__init__.py similarity index 100% rename from src/sparsevllm/tilelang_kernel/mla/__init__.py rename to src/sparsevllm/kernels/tilelang/mla/__init__.py diff --git a/src/sparsevllm/tilelang_kernel/mla/decode.py b/src/sparsevllm/kernels/tilelang/mla/decode.py similarity index 100% rename from src/sparsevllm/tilelang_kernel/mla/decode.py rename to src/sparsevllm/kernels/tilelang/mla/decode.py diff --git a/src/sparsevllm/operators/tilelang_mla.py b/src/sparsevllm/kernels/tilelang/mla/runtime.py similarity index 99% rename from src/sparsevllm/operators/tilelang_mla.py rename to src/sparsevllm/kernels/tilelang/mla/runtime.py index c37a3da9..85152e80 100644 --- a/src/sparsevllm/operators/tilelang_mla.py +++ b/src/sparsevllm/kernels/tilelang/mla/runtime.py @@ -7,13 +7,12 @@ from __future__ import annotations +from collections.abc import Callable from dataclasses import dataclass from importlib import metadata -from typing import Callable import torch - _VALIDATED_TILELANG_VERSION = "0.1.9" _VALIDATED_TVM_FFI_VERSION = "0.1.10" _VALID_SPLITS = (1, 2, 4, 8, 16, 32) @@ -243,7 +242,7 @@ def _bind(self, key: _KernelKey) -> _BoundKernel: raise RuntimeError(reason) # Importing TileLang can initialize its compiler, so keep it behind the # selected provider and outside module import/resolver paths. - from sparsevllm.tilelang_kernel.mla.decode import ( + from sparsevllm.kernels.tilelang.mla.decode import ( build_glm_mla_decode_kernel, ) @@ -465,9 +464,10 @@ def __call__( bound = self._bind(key) self._kernels[key] = bound - from sparsevllm.tilelang_kernel.mla.decode import pad_glm_q_kernel import triton + from sparsevllm.kernels.tilelang.mla.decode import pad_glm_q_kernel + workspace = bound.workspace pad_glm_q_kernel[ (triton.cdiv(batch_size * self.padded_heads * _LATENT_DIM, 256),) diff --git a/src/sparsevllm/triton_kernel/__init__.py b/src/sparsevllm/kernels/triton/__init__.py similarity index 100% rename from src/sparsevllm/triton_kernel/__init__.py rename to src/sparsevllm/kernels/triton/__init__.py diff --git a/src/sparsevllm/triton_kernel/column_parallel_rmsnorm.py b/src/sparsevllm/kernels/triton/column_parallel_rmsnorm.py similarity index 100% rename from src/sparsevllm/triton_kernel/column_parallel_rmsnorm.py rename to src/sparsevllm/kernels/triton/column_parallel_rmsnorm.py diff --git a/src/sparsevllm/triton_kernel/context_flashattention_nopad.py b/src/sparsevllm/kernels/triton/context_flashattention_nopad.py similarity index 100% rename from src/sparsevllm/triton_kernel/context_flashattention_nopad.py rename to src/sparsevllm/kernels/triton/context_flashattention_nopad.py diff --git a/src/sparsevllm/triton_kernel/deltakv_kernels.py b/src/sparsevllm/kernels/triton/deltakv_kernels.py similarity index 100% rename from src/sparsevllm/triton_kernel/deltakv_kernels.py rename to src/sparsevllm/kernels/triton/deltakv_kernels.py diff --git a/src/sparsevllm/triton_kernel/embedding.py b/src/sparsevllm/kernels/triton/embedding.py similarity index 100% rename from src/sparsevllm/triton_kernel/embedding.py rename to src/sparsevllm/kernels/triton/embedding.py diff --git a/src/sparsevllm/triton_kernel/flash_decoding.py b/src/sparsevllm/kernels/triton/flash_decoding.py similarity index 100% rename from src/sparsevllm/triton_kernel/flash_decoding.py rename to src/sparsevllm/kernels/triton/flash_decoding.py diff --git a/src/sparsevllm/triton_kernel/flash_decoding_stage1.py b/src/sparsevllm/kernels/triton/flash_decoding_stage1.py similarity index 100% rename from src/sparsevllm/triton_kernel/flash_decoding_stage1.py rename to src/sparsevllm/kernels/triton/flash_decoding_stage1.py diff --git a/src/sparsevllm/triton_kernel/flash_decoding_stage2.py b/src/sparsevllm/kernels/triton/flash_decoding_stage2.py similarity index 100% rename from src/sparsevllm/triton_kernel/flash_decoding_stage2.py rename to src/sparsevllm/kernels/triton/flash_decoding_stage2.py diff --git a/src/sparsevllm/triton_kernel/fp8_blockwise.py b/src/sparsevllm/kernels/triton/fp8_blockwise.py similarity index 100% rename from src/sparsevllm/triton_kernel/fp8_blockwise.py rename to src/sparsevllm/kernels/triton/fp8_blockwise.py diff --git a/src/sparsevllm/triton_kernel/gate_up_swiglu.py b/src/sparsevllm/kernels/triton/gate_up_swiglu.py similarity index 100% rename from src/sparsevllm/triton_kernel/gate_up_swiglu.py rename to src/sparsevllm/kernels/triton/gate_up_swiglu.py diff --git a/src/sparsevllm/triton_kernel/gqa_decode_flashattention_nopad.py b/src/sparsevllm/kernels/triton/gqa_decode_flashattention_nopad.py similarity index 100% rename from src/sparsevllm/triton_kernel/gqa_decode_flashattention_nopad.py rename to src/sparsevllm/kernels/triton/gqa_decode_flashattention_nopad.py diff --git a/src/sparsevllm/triton_kernel/gqa_flash_decoding.py b/src/sparsevllm/kernels/triton/gqa_flash_decoding.py similarity index 100% rename from src/sparsevllm/triton_kernel/gqa_flash_decoding.py rename to src/sparsevllm/kernels/triton/gqa_flash_decoding.py diff --git a/src/sparsevllm/triton_kernel/gqa_flash_decoding_stage1.py b/src/sparsevllm/kernels/triton/gqa_flash_decoding_stage1.py similarity index 100% rename from src/sparsevllm/triton_kernel/gqa_flash_decoding_stage1.py rename to src/sparsevllm/kernels/triton/gqa_flash_decoding_stage1.py diff --git a/src/sparsevllm/triton_kernel/gqa_flash_decoding_stage2.py b/src/sparsevllm/kernels/triton/gqa_flash_decoding_stage2.py similarity index 100% rename from src/sparsevllm/triton_kernel/gqa_flash_decoding_stage2.py rename to src/sparsevllm/kernels/triton/gqa_flash_decoding_stage2.py diff --git a/src/sparsevllm/triton_kernel/minimax_m2_moe.py b/src/sparsevllm/kernels/triton/minimax_m2_moe.py similarity index 99% rename from src/sparsevllm/triton_kernel/minimax_m2_moe.py rename to src/sparsevllm/kernels/triton/minimax_m2_moe.py index 7ee34751..a4ed4625 100644 --- a/src/sparsevllm/triton_kernel/minimax_m2_moe.py +++ b/src/sparsevllm/kernels/triton/minimax_m2_moe.py @@ -4,7 +4,7 @@ import triton import triton.language as tl -from sparsevllm.triton_kernel.moe import ( +from sparsevllm.kernels.triton.moe import ( _prepare_expert_assignment, _routed_fp8_gemm, moe_sum, diff --git a/src/sparsevllm/triton_kernel/minimax_m2_router.py b/src/sparsevllm/kernels/triton/minimax_m2_router.py similarity index 98% rename from src/sparsevllm/triton_kernel/minimax_m2_router.py rename to src/sparsevllm/kernels/triton/minimax_m2_router.py index 48c9004a..3c613d74 100644 --- a/src/sparsevllm/triton_kernel/minimax_m2_router.py +++ b/src/sparsevllm/kernels/triton/minimax_m2_router.py @@ -3,7 +3,7 @@ import torch import torch.nn.functional as F -from sparsevllm.triton_kernel.moe_biased_sigmoid import ( +from sparsevllm.kernels.triton.moe_biased_sigmoid import ( topk_biased_sigmoid as _topk_biased_sigmoid, ) diff --git a/src/sparsevllm/triton_kernel/mla/LICENSE.lightllm b/src/sparsevllm/kernels/triton/mla/LICENSE.lightllm similarity index 100% rename from src/sparsevllm/triton_kernel/mla/LICENSE.lightllm rename to src/sparsevllm/kernels/triton/mla/LICENSE.lightllm diff --git a/src/sparsevllm/triton_kernel/mla/README.md b/src/sparsevllm/kernels/triton/mla/README.md similarity index 100% rename from src/sparsevllm/triton_kernel/mla/README.md rename to src/sparsevllm/kernels/triton/mla/README.md diff --git a/src/sparsevllm/triton_kernel/mla/__init__.py b/src/sparsevllm/kernels/triton/mla/__init__.py similarity index 100% rename from src/sparsevllm/triton_kernel/mla/__init__.py rename to src/sparsevllm/kernels/triton/mla/__init__.py diff --git a/src/sparsevllm/triton_kernel/mla/copy_latent.py b/src/sparsevllm/kernels/triton/mla/copy_latent.py similarity index 100% rename from src/sparsevllm/triton_kernel/mla/copy_latent.py rename to src/sparsevllm/kernels/triton/mla/copy_latent.py diff --git a/src/sparsevllm/triton_kernel/mla/decode_schedule.py b/src/sparsevllm/kernels/triton/mla/decode_schedule.py similarity index 100% rename from src/sparsevllm/triton_kernel/mla/decode_schedule.py rename to src/sparsevllm/kernels/triton/mla/decode_schedule.py diff --git a/src/sparsevllm/triton_kernel/mla/decode_stage1.py b/src/sparsevllm/kernels/triton/mla/decode_stage1.py similarity index 100% rename from src/sparsevllm/triton_kernel/mla/decode_stage1.py rename to src/sparsevllm/kernels/triton/mla/decode_stage1.py diff --git a/src/sparsevllm/triton_kernel/mla/decode_stage2.py b/src/sparsevllm/kernels/triton/mla/decode_stage2.py similarity index 100% rename from src/sparsevllm/triton_kernel/mla/decode_stage2.py rename to src/sparsevllm/kernels/triton/mla/decode_stage2.py diff --git a/src/sparsevllm/triton_kernel/mla/gather_latent.py b/src/sparsevllm/kernels/triton/mla/gather_latent.py similarity index 100% rename from src/sparsevllm/triton_kernel/mla/gather_latent.py rename to src/sparsevllm/kernels/triton/mla/gather_latent.py diff --git a/src/sparsevllm/triton_kernel/moe.py b/src/sparsevllm/kernels/triton/moe.py similarity index 99% rename from src/sparsevllm/triton_kernel/moe.py rename to src/sparsevllm/kernels/triton/moe.py index 7b83cd9f..187d7ce0 100644 --- a/src/sparsevllm/triton_kernel/moe.py +++ b/src/sparsevllm/kernels/triton/moe.py @@ -1,33 +1,23 @@ from __future__ import annotations from collections.abc import Callable -from dataclasses import dataclass import torch import triton import triton.language as tl -from sparsevllm.triton_kernel.silu_and_mul import silu_and_mul_fwd -from sparsevllm.triton_kernel.moe_config import ( +from sparsevllm.kernels.moe import MoeAlignment +from sparsevllm.kernels.triton.moe_config import ( MoeGemmConfig, device_info, resolve_fp8_routed_gemm_config, resolve_moe_gemm_config, ) - +from sparsevllm.kernels.triton.silu_and_mul import silu_and_mul_fwd _SUPPORTED_DTYPES = (torch.bfloat16, torch.float16) -@dataclass(frozen=True) -class MoeAlignment: - sorted_token_ids: torch.Tensor | None - expert_ids: torch.Tensor - num_tokens_post_padded: torch.Tensor - block_size: int - naive: bool - - @triton.jit def _localize_expert_ids_kernel( input_ids_ptr, diff --git a/src/sparsevllm/triton_kernel/moe_biased_sigmoid.py b/src/sparsevllm/kernels/triton/moe_biased_sigmoid.py similarity index 100% rename from src/sparsevllm/triton_kernel/moe_biased_sigmoid.py rename to src/sparsevllm/kernels/triton/moe_biased_sigmoid.py diff --git a/src/sparsevllm/triton_kernel/moe_config.py b/src/sparsevllm/kernels/triton/moe_config.py similarity index 100% rename from src/sparsevllm/triton_kernel/moe_config.py rename to src/sparsevllm/kernels/triton/moe_config.py diff --git a/src/sparsevllm/triton_kernel/moe_topk.py b/src/sparsevllm/kernels/triton/moe_topk.py similarity index 100% rename from src/sparsevllm/triton_kernel/moe_topk.py rename to src/sparsevllm/kernels/triton/moe_topk.py diff --git a/src/sparsevllm/triton_kernel/omnikv_fused.py b/src/sparsevllm/kernels/triton/omnikv_fused.py similarity index 100% rename from src/sparsevllm/triton_kernel/omnikv_fused.py rename to src/sparsevllm/kernels/triton/omnikv_fused.py diff --git a/src/sparsevllm/triton_kernel/ppl_fp16_flash_decoding.py b/src/sparsevllm/kernels/triton/ppl_fp16_flash_decoding.py similarity index 100% rename from src/sparsevllm/triton_kernel/ppl_fp16_flash_decoding.py rename to src/sparsevllm/kernels/triton/ppl_fp16_flash_decoding.py diff --git a/src/sparsevllm/triton_kernel/ppl_int4kv_copy_kv.py b/src/sparsevllm/kernels/triton/ppl_int4kv_copy_kv.py similarity index 100% rename from src/sparsevllm/triton_kernel/ppl_int4kv_copy_kv.py rename to src/sparsevllm/kernels/triton/ppl_int4kv_copy_kv.py diff --git a/src/sparsevllm/triton_kernel/ppl_int4kv_flash_decoding.py b/src/sparsevllm/kernels/triton/ppl_int4kv_flash_decoding.py similarity index 100% rename from src/sparsevllm/triton_kernel/ppl_int4kv_flash_decoding.py rename to src/sparsevllm/kernels/triton/ppl_int4kv_flash_decoding.py diff --git a/src/sparsevllm/triton_kernel/ppl_int8kv_flash_decoding.py b/src/sparsevllm/kernels/triton/ppl_int8kv_flash_decoding.py similarity index 100% rename from src/sparsevllm/triton_kernel/ppl_int8kv_flash_decoding.py rename to src/sparsevllm/kernels/triton/ppl_int8kv_flash_decoding.py diff --git a/src/sparsevllm/triton_kernel/ppl_quant_copy_kv.py b/src/sparsevllm/kernels/triton/ppl_quant_copy_kv.py similarity index 100% rename from src/sparsevllm/triton_kernel/ppl_quant_copy_kv.py rename to src/sparsevllm/kernels/triton/ppl_quant_copy_kv.py diff --git a/src/sparsevllm/triton_kernel/prefill_score.py b/src/sparsevllm/kernels/triton/prefill_score.py similarity index 100% rename from src/sparsevllm/triton_kernel/prefill_score.py rename to src/sparsevllm/kernels/triton/prefill_score.py diff --git a/src/sparsevllm/triton_kernel/quant.py b/src/sparsevllm/kernels/triton/quant.py similarity index 100% rename from src/sparsevllm/triton_kernel/quant.py rename to src/sparsevllm/kernels/triton/quant.py diff --git a/src/sparsevllm/triton_kernel/qwen3_5/README.md b/src/sparsevllm/kernels/triton/qwen3_5/README.md similarity index 100% rename from src/sparsevllm/triton_kernel/qwen3_5/README.md rename to src/sparsevllm/kernels/triton/qwen3_5/README.md diff --git a/src/sparsevllm/triton_kernel/qwen3_5/__init__.py b/src/sparsevllm/kernels/triton/qwen3_5/__init__.py similarity index 100% rename from src/sparsevllm/triton_kernel/qwen3_5/__init__.py rename to src/sparsevllm/kernels/triton/qwen3_5/__init__.py diff --git a/src/sparsevllm/triton_kernel/qwen3_5/autotuner.py b/src/sparsevllm/kernels/triton/qwen3_5/autotuner.py similarity index 100% rename from src/sparsevllm/triton_kernel/qwen3_5/autotuner.py rename to src/sparsevllm/kernels/triton/qwen3_5/autotuner.py diff --git a/src/sparsevllm/triton_kernel/qwen3_5/causal_conv1d.py b/src/sparsevllm/kernels/triton/qwen3_5/causal_conv1d.py similarity index 100% rename from src/sparsevllm/triton_kernel/qwen3_5/causal_conv1d.py rename to src/sparsevllm/kernels/triton/qwen3_5/causal_conv1d.py diff --git a/src/sparsevllm/triton_kernel/qwen3_5/fla/__init__.py b/src/sparsevllm/kernels/triton/qwen3_5/fla/__init__.py similarity index 100% rename from src/sparsevllm/triton_kernel/qwen3_5/fla/__init__.py rename to src/sparsevllm/kernels/triton/qwen3_5/fla/__init__.py diff --git a/src/sparsevllm/triton_kernel/qwen3_5/fla/ops/__init__.py b/src/sparsevllm/kernels/triton/qwen3_5/fla/ops/__init__.py similarity index 100% rename from src/sparsevllm/triton_kernel/qwen3_5/fla/ops/__init__.py rename to src/sparsevllm/kernels/triton/qwen3_5/fla/ops/__init__.py diff --git a/src/sparsevllm/triton_kernel/qwen3_5/fla/ops/chunk.py b/src/sparsevllm/kernels/triton/qwen3_5/fla/ops/chunk.py similarity index 100% rename from src/sparsevllm/triton_kernel/qwen3_5/fla/ops/chunk.py rename to src/sparsevllm/kernels/triton/qwen3_5/fla/ops/chunk.py diff --git a/src/sparsevllm/triton_kernel/qwen3_5/fla/ops/chunk_delta_h.py b/src/sparsevllm/kernels/triton/qwen3_5/fla/ops/chunk_delta_h.py similarity index 99% rename from src/sparsevllm/triton_kernel/qwen3_5/fla/ops/chunk_delta_h.py rename to src/sparsevllm/kernels/triton/qwen3_5/fla/ops/chunk_delta_h.py index 4506c65e..3e6fec20 100644 --- a/src/sparsevllm/triton_kernel/qwen3_5/fla/ops/chunk_delta_h.py +++ b/src/sparsevllm/kernels/triton/qwen3_5/fla/ops/chunk_delta_h.py @@ -15,7 +15,7 @@ from .index import prepare_chunk_indices, prepare_chunk_offsets from .op import exp, safe_exp -from sparsevllm.triton_kernel.qwen3_5.autotuner import autotune +from sparsevllm.kernels.triton.qwen3_5.autotuner import autotune NUM_WARPS = [2, 4, 8, 16] diff --git a/src/sparsevllm/triton_kernel/qwen3_5/fla/ops/chunk_o.py b/src/sparsevllm/kernels/triton/qwen3_5/fla/ops/chunk_o.py similarity index 98% rename from src/sparsevllm/triton_kernel/qwen3_5/fla/ops/chunk_o.py rename to src/sparsevllm/kernels/triton/qwen3_5/fla/ops/chunk_o.py index 8b196dc7..7d5af2be 100644 --- a/src/sparsevllm/triton_kernel/qwen3_5/fla/ops/chunk_o.py +++ b/src/sparsevllm/kernels/triton/qwen3_5/fla/ops/chunk_o.py @@ -18,7 +18,7 @@ from .index import prepare_chunk_indices from .op import exp, safe_exp from .utils import FLA_GDN_FIX_BT, check_shared_mem, is_nvidia_hopper -from sparsevllm.triton_kernel.qwen3_5.autotuner import autotune +from sparsevllm.kernels.triton.qwen3_5.autotuner import autotune BKV_LIST = [64, 128] if check_shared_mem() else [32, 64] NUM_WARPS = [2, 4] if is_nvidia_hopper else [2, 4, 8] diff --git a/src/sparsevllm/triton_kernel/qwen3_5/fla/ops/chunk_scaled_dot_kkt.py b/src/sparsevllm/kernels/triton/qwen3_5/fla/ops/chunk_scaled_dot_kkt.py similarity index 98% rename from src/sparsevllm/triton_kernel/qwen3_5/fla/ops/chunk_scaled_dot_kkt.py rename to src/sparsevllm/kernels/triton/qwen3_5/fla/ops/chunk_scaled_dot_kkt.py index 5c17b31f..ac7c2e27 100644 --- a/src/sparsevllm/triton_kernel/qwen3_5/fla/ops/chunk_scaled_dot_kkt.py +++ b/src/sparsevllm/kernels/triton/qwen3_5/fla/ops/chunk_scaled_dot_kkt.py @@ -15,7 +15,7 @@ from .index import prepare_chunk_indices from .op import safe_exp -from sparsevllm.triton_kernel.qwen3_5.autotuner import autotune +from sparsevllm.kernels.triton.qwen3_5.autotuner import autotune @triton.heuristics( diff --git a/src/sparsevllm/triton_kernel/qwen3_5/fla/ops/cumsum.py b/src/sparsevllm/kernels/triton/qwen3_5/fla/ops/cumsum.py similarity index 99% rename from src/sparsevllm/triton_kernel/qwen3_5/fla/ops/cumsum.py rename to src/sparsevllm/kernels/triton/qwen3_5/fla/ops/cumsum.py index fd680025..cedbb789 100644 --- a/src/sparsevllm/triton_kernel/qwen3_5/fla/ops/cumsum.py +++ b/src/sparsevllm/kernels/triton/qwen3_5/fla/ops/cumsum.py @@ -14,7 +14,7 @@ from .index import prepare_chunk_indices from .utils import check_shared_mem, input_guard -from sparsevllm.triton_kernel.qwen3_5.autotuner import autotune +from sparsevllm.kernels.triton.qwen3_5.autotuner import autotune BS_LIST = [32, 64] if check_shared_mem() else [16, 32] diff --git a/src/sparsevllm/triton_kernel/qwen3_5/fla/ops/fused_recurrent.py b/src/sparsevllm/kernels/triton/qwen3_5/fla/ops/fused_recurrent.py similarity index 100% rename from src/sparsevllm/triton_kernel/qwen3_5/fla/ops/fused_recurrent.py rename to src/sparsevllm/kernels/triton/qwen3_5/fla/ops/fused_recurrent.py diff --git a/src/sparsevllm/triton_kernel/qwen3_5/fla/ops/index.py b/src/sparsevllm/kernels/triton/qwen3_5/fla/ops/index.py similarity index 100% rename from src/sparsevllm/triton_kernel/qwen3_5/fla/ops/index.py rename to src/sparsevllm/kernels/triton/qwen3_5/fla/ops/index.py diff --git a/src/sparsevllm/triton_kernel/qwen3_5/fla/ops/l2norm.py b/src/sparsevllm/kernels/triton/qwen3_5/fla/ops/l2norm.py similarity index 98% rename from src/sparsevllm/triton_kernel/qwen3_5/fla/ops/l2norm.py rename to src/sparsevllm/kernels/triton/qwen3_5/fla/ops/l2norm.py index 2cb42818..59ba3161 100644 --- a/src/sparsevllm/triton_kernel/qwen3_5/fla/ops/l2norm.py +++ b/src/sparsevllm/kernels/triton/qwen3_5/fla/ops/l2norm.py @@ -13,7 +13,7 @@ import triton import triton.language as tl -from sparsevllm.triton_kernel.qwen3_5.autotuner import autotune +from sparsevllm.kernels.triton.qwen3_5.autotuner import autotune BT_LIST = [8, 16, 32, 64, 128] diff --git a/src/sparsevllm/triton_kernel/qwen3_5/fla/ops/op.py b/src/sparsevllm/kernels/triton/qwen3_5/fla/ops/op.py similarity index 100% rename from src/sparsevllm/triton_kernel/qwen3_5/fla/ops/op.py rename to src/sparsevllm/kernels/triton/qwen3_5/fla/ops/op.py diff --git a/src/sparsevllm/triton_kernel/qwen3_5/fla/ops/solve_tril.py b/src/sparsevllm/kernels/triton/qwen3_5/fla/ops/solve_tril.py similarity index 100% rename from src/sparsevllm/triton_kernel/qwen3_5/fla/ops/solve_tril.py rename to src/sparsevllm/kernels/triton/qwen3_5/fla/ops/solve_tril.py diff --git a/src/sparsevllm/triton_kernel/qwen3_5/fla/ops/utils.py b/src/sparsevllm/kernels/triton/qwen3_5/fla/ops/utils.py similarity index 100% rename from src/sparsevllm/triton_kernel/qwen3_5/fla/ops/utils.py rename to src/sparsevllm/kernels/triton/qwen3_5/fla/ops/utils.py diff --git a/src/sparsevllm/triton_kernel/qwen3_5/fla/ops/wy_fast.py b/src/sparsevllm/kernels/triton/qwen3_5/fla/ops/wy_fast.py similarity index 100% rename from src/sparsevllm/triton_kernel/qwen3_5/fla/ops/wy_fast.py rename to src/sparsevllm/kernels/triton/qwen3_5/fla/ops/wy_fast.py diff --git a/src/sparsevllm/triton_kernel/qwen3_5/fused_gdn_gating.py b/src/sparsevllm/kernels/triton/qwen3_5/fused_gdn_gating.py similarity index 97% rename from src/sparsevllm/triton_kernel/qwen3_5/fused_gdn_gating.py rename to src/sparsevllm/kernels/triton/qwen3_5/fused_gdn_gating.py index 91fb7b8a..bf11f307 100644 --- a/src/sparsevllm/triton_kernel/qwen3_5/fused_gdn_gating.py +++ b/src/sparsevllm/kernels/triton/qwen3_5/fused_gdn_gating.py @@ -5,7 +5,7 @@ import triton import triton.language as tl -from sparsevllm.triton_kernel.qwen3_5.autotuner import autotune +from sparsevllm.kernels.triton.qwen3_5.autotuner import autotune # g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias) diff --git a/src/sparsevllm/triton_kernel/qwen3_5/gated_rmsnorm.py b/src/sparsevllm/kernels/triton/qwen3_5/gated_rmsnorm.py similarity index 98% rename from src/sparsevllm/triton_kernel/qwen3_5/gated_rmsnorm.py rename to src/sparsevllm/kernels/triton/qwen3_5/gated_rmsnorm.py index 65f51c0d..13d1d565 100644 --- a/src/sparsevllm/triton_kernel/qwen3_5/gated_rmsnorm.py +++ b/src/sparsevllm/kernels/triton/qwen3_5/gated_rmsnorm.py @@ -1,7 +1,7 @@ import triton import triton.language as tl import torch -from sparsevllm.triton_kernel.qwen3_5.autotuner import autotune +from sparsevllm.kernels.triton.qwen3_5.autotuner import autotune @triton.heuristics( diff --git a/src/sparsevllm/triton_kernel/qwen3_5/gated_shared_add.py b/src/sparsevllm/kernels/triton/qwen3_5/gated_shared_add.py similarity index 100% rename from src/sparsevllm/triton_kernel/qwen3_5/gated_shared_add.py rename to src/sparsevllm/kernels/triton/qwen3_5/gated_shared_add.py diff --git a/src/sparsevllm/triton_kernel/qwen3_5/gdn_decode_pack.py b/src/sparsevllm/kernels/triton/qwen3_5/gdn_decode_pack.py similarity index 100% rename from src/sparsevllm/triton_kernel/qwen3_5/gdn_decode_pack.py rename to src/sparsevllm/kernels/triton/qwen3_5/gdn_decode_pack.py diff --git a/src/sparsevllm/triton_kernel/rmsnorm.py b/src/sparsevllm/kernels/triton/rmsnorm.py similarity index 100% rename from src/sparsevllm/triton_kernel/rmsnorm.py rename to src/sparsevllm/kernels/triton/rmsnorm.py diff --git a/src/sparsevllm/triton_kernel/rotary_emb.py b/src/sparsevllm/kernels/triton/rotary_emb.py similarity index 100% rename from src/sparsevllm/triton_kernel/rotary_emb.py rename to src/sparsevllm/kernels/triton/rotary_emb.py diff --git a/src/sparsevllm/triton_kernel/silu_and_mul.py b/src/sparsevllm/kernels/triton/silu_and_mul.py similarity index 100% rename from src/sparsevllm/triton_kernel/silu_and_mul.py rename to src/sparsevllm/kernels/triton/silu_and_mul.py diff --git a/src/sparsevllm/triton_kernel/splitfuse_context_flashattention_nopad.py b/src/sparsevllm/kernels/triton/splitfuse_context_flashattention_nopad.py similarity index 100% rename from src/sparsevllm/triton_kernel/splitfuse_context_flashattention_nopad.py rename to src/sparsevllm/kernels/triton/splitfuse_context_flashattention_nopad.py diff --git a/src/sparsevllm/triton_kernel/store_kvcache.py b/src/sparsevllm/kernels/triton/store_kvcache.py similarity index 100% rename from src/sparsevllm/triton_kernel/store_kvcache.py rename to src/sparsevllm/kernels/triton/store_kvcache.py diff --git a/src/sparsevllm/triton_kernel/token_attention_nopad_att1.py b/src/sparsevllm/kernels/triton/token_attention_nopad_att1.py similarity index 100% rename from src/sparsevllm/triton_kernel/token_attention_nopad_att1.py rename to src/sparsevllm/kernels/triton/token_attention_nopad_att1.py diff --git a/src/sparsevllm/triton_kernel/token_attention_nopad_reduceV.py b/src/sparsevllm/kernels/triton/token_attention_nopad_reduceV.py similarity index 100% rename from src/sparsevllm/triton_kernel/token_attention_nopad_reduceV.py rename to src/sparsevllm/kernels/triton/token_attention_nopad_reduceV.py diff --git a/src/sparsevllm/triton_kernel/token_attention_nopad_softmax.py b/src/sparsevllm/kernels/triton/token_attention_nopad_softmax.py similarity index 100% rename from src/sparsevllm/triton_kernel/token_attention_nopad_softmax.py rename to src/sparsevllm/kernels/triton/token_attention_nopad_softmax.py diff --git a/src/sparsevllm/triton_kernel/token_attention_softmax_and_reducev.py b/src/sparsevllm/kernels/triton/token_attention_softmax_and_reducev.py similarity index 100% rename from src/sparsevllm/triton_kernel/token_attention_softmax_and_reducev.py rename to src/sparsevllm/kernels/triton/token_attention_softmax_and_reducev.py diff --git a/src/sparsevllm/layers/attention_backend.py b/src/sparsevllm/layers/attention_backend.py index 5d13e474..02a93372 100644 --- a/src/sparsevllm/layers/attention_backend.py +++ b/src/sparsevllm/layers/attention_backend.py @@ -9,12 +9,12 @@ ) from sparsevllm.operators.registry import record_operator_binding from sparsevllm.utils.context import get_context -from sparsevllm.triton_kernel.context_flashattention_nopad import context_attention_fwd -from sparsevllm.triton_kernel.flash_decoding_stage1 import flash_decode_stage1 as mha_flash_decode_stage1 -from sparsevllm.triton_kernel.flash_decoding_stage1 import flash_decode_stage1_with_score as mha_flash_decode_stage1_with_score -from sparsevllm.triton_kernel.flash_decoding_stage2 import flash_decode_stage2 -from sparsevllm.triton_kernel.gqa_flash_decoding_stage1 import flash_decode_stage1 as gqa_flash_decode_stage1 -from sparsevllm.triton_kernel.gqa_flash_decoding_stage1 import flash_decode_stage1_with_score as gqa_flash_decode_stage1_with_score +from sparsevllm.kernels.triton.context_flashattention_nopad import context_attention_fwd +from sparsevllm.kernels.triton.flash_decoding_stage1 import flash_decode_stage1 as mha_flash_decode_stage1 +from sparsevllm.kernels.triton.flash_decoding_stage1 import flash_decode_stage1_with_score as mha_flash_decode_stage1_with_score +from sparsevllm.kernels.triton.flash_decoding_stage2 import flash_decode_stage2 +from sparsevllm.kernels.triton.gqa_flash_decoding_stage1 import flash_decode_stage1 as gqa_flash_decode_stage1 +from sparsevllm.kernels.triton.gqa_flash_decoding_stage1 import flash_decode_stage1_with_score as gqa_flash_decode_stage1_with_score from sparsevllm.utils.log import log_once from sparsevllm.utils.profiler import profiler @@ -418,7 +418,7 @@ def _run_full_layer_kivi_decode_stage1( backend_metadata = payload.metadata if backend_metadata is None: raise RuntimeError("full_layer_kivi decode view is missing metadata.") - from sparsevllm.triton_kernel.deltakv_kernels import full_layer_kivi_flash_decode_stage1 + from sparsevllm.kernels.triton.deltakv_kernels import full_layer_kivi_flash_decode_stage1 full_layer_kivi_flash_decode_stage1( q=q, diff --git a/src/sparsevllm/layers/layernorm.py b/src/sparsevllm/layers/layernorm.py index a7afb418..c956fadb 100755 --- a/src/sparsevllm/layers/layernorm.py +++ b/src/sparsevllm/layers/layernorm.py @@ -59,7 +59,7 @@ def run_fused_add_rmsnorm( def _load_triton_ops(*, zero_centered_weight: bool) -> _RMSNormOps: - from sparsevllm.triton_kernel.rmsnorm import ( + from sparsevllm.kernels.triton.rmsnorm import ( fused_add_rmsnorm_forward, rmsnorm_forward, ) @@ -223,7 +223,7 @@ def forward_pair( if self.parallel_context is not other_norm.parallel_context: raise ValueError("Paired column-parallel RMSNorms must share a context.") if x.is_cuda or other.is_cuda: - from sparsevllm.triton_kernel.column_parallel_rmsnorm import ( + from sparsevllm.kernels.triton.column_parallel_rmsnorm import ( paired_rms_apply, paired_square_sums, ) diff --git a/src/sparsevllm/layers/mla_attention.py b/src/sparsevllm/layers/mla_attention.py index 04b14d5c..2a7eac5c 100644 --- a/src/sparsevllm/layers/mla_attention.py +++ b/src/sparsevllm/layers/mla_attention.py @@ -19,7 +19,7 @@ MlaAttentionProvider, resolve_mla_attention_provider, ) -from sparsevllm.triton_kernel.mla import ( +from sparsevllm.kernels.triton.mla import ( gather_latent_history, validate_gather_metadata, ) diff --git a/src/sparsevllm/models/glm4_moe_lite.py b/src/sparsevllm/models/glm4_moe_lite.py index 4670db15..123379c1 100644 --- a/src/sparsevllm/models/glm4_moe_lite.py +++ b/src/sparsevllm/models/glm4_moe_lite.py @@ -628,7 +628,7 @@ def forward_routed_and_shared( ) -> torch.Tensor: if self.shared_expert_id is None: raise RuntimeError("Packed shared expert is not enabled.") - from sparsevllm.triton_kernel.moe import append_shared_expert_route + from sparsevllm.kernels.triton.moe import append_shared_expert_route fused_ids, fused_weights = append_shared_expert_route( topk_ids, diff --git a/src/sparsevllm/models/minimax_m2.py b/src/sparsevllm/models/minimax_m2.py index 1e47b087..1efcd448 100644 --- a/src/sparsevllm/models/minimax_m2.py +++ b/src/sparsevllm/models/minimax_m2.py @@ -199,7 +199,7 @@ def __init__(self, config) -> None: self.e_score_correction_bias = nn.Parameter( torch.empty(self.num_experts, dtype=torch.float32) ) - from sparsevllm.triton_kernel.minimax_m2_router import ( + from sparsevllm.kernels.triton.minimax_m2_router import ( topk_biased_sigmoid, ) diff --git a/src/sparsevllm/models/qwen3_5.py b/src/sparsevllm/models/qwen3_5.py index 2ba37904..71a23c4d 100644 --- a/src/sparsevllm/models/qwen3_5.py +++ b/src/sparsevllm/models/qwen3_5.py @@ -24,11 +24,11 @@ from sparsevllm.utils.context import get_context from sparsevllm.utils.weight_target import WeightTarget from sparsevllm.engine.recurrent_state_manager import RecurrentStateSpec, RecurrentTensorSpec -from sparsevllm.triton_kernel.qwen3_5.causal_conv1d import causal_conv1d_fn -from sparsevllm.triton_kernel.qwen3_5.fused_gdn_gating import fused_gdn_gating -from sparsevllm.triton_kernel.qwen3_5.gated_rmsnorm import gated_rmsnorm_forward -from sparsevllm.triton_kernel.qwen3_5.gdn_decode_pack import conv_pack_gdn_decode_inputs -from sparsevllm.triton_kernel.qwen3_5.fla.ops import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule +from sparsevllm.kernels.triton.qwen3_5.causal_conv1d import causal_conv1d_fn +from sparsevllm.kernels.triton.qwen3_5.fused_gdn_gating import fused_gdn_gating +from sparsevllm.kernels.triton.qwen3_5.gated_rmsnorm import gated_rmsnorm_forward +from sparsevllm.kernels.triton.qwen3_5.gdn_decode_pack import conv_pack_gdn_decode_inputs +from sparsevllm.kernels.triton.qwen3_5.fla.ops import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule def _get_rope_theta(config) -> float: diff --git a/src/sparsevllm/models/qwen3_moe.py b/src/sparsevllm/models/qwen3_moe.py index eefe0153..a438169c 100644 --- a/src/sparsevllm/models/qwen3_moe.py +++ b/src/sparsevllm/models/qwen3_moe.py @@ -38,7 +38,7 @@ def __init__(self, config: Qwen3MoeConfig) -> None: self.num_experts = int(config.num_experts) self.top_k = int(config.num_experts_per_tok) self.norm_topk_prob = bool(config.norm_topk_prob) - from sparsevllm.triton_kernel.moe_topk import topk_softmax + from sparsevllm.kernels.triton.moe_topk import topk_softmax self.topk_impl = topk_softmax self.weight = nn.Parameter(torch.empty(self.num_experts, self.hidden_size)) diff --git a/src/sparsevllm/operators/activation.py b/src/sparsevllm/operators/activation.py index dea72f33..072c9c00 100644 --- a/src/sparsevllm/operators/activation.py +++ b/src/sparsevllm/operators/activation.py @@ -85,7 +85,7 @@ def __call__(self, x: torch.Tensor) -> torch.Tensor: _validate_bound_input(x, self.spec) if not x.is_cuda: raise ValueError("Triton SiluAndMul provider requires a CUDA input.") - from sparsevllm.triton_kernel.silu_and_mul import silu_and_mul_fwd + from sparsevllm.kernels.triton.silu_and_mul import silu_and_mul_fwd return silu_and_mul_fwd(x) diff --git a/src/sparsevllm/operators/fp8_linear.py b/src/sparsevllm/operators/fp8_linear.py index 0ba23627..39425eff 100644 --- a/src/sparsevllm/operators/fp8_linear.py +++ b/src/sparsevllm/operators/fp8_linear.py @@ -125,7 +125,7 @@ def supports(cls, spec: Fp8LinearSpec, caps: DeviceCaps) -> SupportResult: return SupportResult.yes() def __call__(self, x, weight, weight_scale_inv, bias=None): - from sparsevllm.triton_kernel.fp8_blockwise import fp8_blockwise_matmul + from sparsevllm.kernels.triton.fp8_blockwise import fp8_blockwise_matmul original_shape = x.shape[:-1] output = fp8_blockwise_matmul( diff --git a/src/sparsevllm/operators/gate_up_swiglu.py b/src/sparsevllm/operators/gate_up_swiglu.py index c4ac8c6a..b2ab9739 100644 --- a/src/sparsevllm/operators/gate_up_swiglu.py +++ b/src/sparsevllm/operators/gate_up_swiglu.py @@ -126,7 +126,7 @@ def run( ) -> torch.Tensor: if inputs.shape[0] != 1: return super().run(spec, inputs, projection) - from sparsevllm.triton_kernel.gate_up_swiglu import h20_gate_up_swiglu + from sparsevllm.kernels.triton.gate_up_swiglu import h20_gate_up_swiglu return h20_gate_up_swiglu(inputs, projection.weight) diff --git a/src/sparsevllm/operators/gated_shared_add.py b/src/sparsevllm/operators/gated_shared_add.py index df9de374..43afba98 100644 --- a/src/sparsevllm/operators/gated_shared_add.py +++ b/src/sparsevllm/operators/gated_shared_add.py @@ -37,7 +37,7 @@ def gated_shared_add( ): raise ValueError("gated_shared_add requires contiguous hidden dimensions.") - from sparsevllm.triton_kernel.qwen3_5.gated_shared_add import ( + from sparsevllm.kernels.triton.qwen3_5.gated_shared_add import ( triton_gated_shared_add, ) diff --git a/src/sparsevllm/operators/mla_attention.py b/src/sparsevllm/operators/mla_attention.py index 5694c6dd..4ad1c96d 100644 --- a/src/sparsevllm/operators/mla_attention.py +++ b/src/sparsevllm/operators/mla_attention.py @@ -11,18 +11,12 @@ MlaLatentPayload, PrefillComputeView, ) -from sparsevllm.operators.registry import ( - OpRegistry, - OpResolver, - SupportResult, -) -from sparsevllm.operators.sgl_fa3 import SglFa3DecodeKernel, sgl_fa3_support -from sparsevllm.operators.tilelang_mla import ( +from sparsevllm.kernels.external.sgl.fa3 import SglFa3DecodeKernel, sgl_fa3_support +from sparsevllm.kernels.tilelang.mla.runtime import ( TileMlaDecodeKernel, tilelang_mla_support, ) -from sparsevllm.platforms.interface import DeviceCaps, PlatformEnum -from sparsevllm.triton_kernel.mla import ( +from sparsevllm.kernels.triton.mla import ( DEFAULT_GLM_MLA_DECODE_CONFIG, GLM_MLA_MAX_WORKSPACE_CONFIG, MlaDecodeLaunchConfig, @@ -31,7 +25,12 @@ select_glm_mla_decode_config, validate_mla_decode_metadata, ) - +from sparsevllm.operators.registry import ( + OpRegistry, + OpResolver, + SupportResult, +) +from sparsevllm.platforms.interface import DeviceCaps, PlatformEnum _GLM_MLA_NUM_Q_HEADS = 20 _GLM_MLA_KV_LORA_RANK = 512 diff --git a/src/sparsevllm/operators/moe.py b/src/sparsevllm/operators/moe.py index 1804874e..2a1148c6 100644 --- a/src/sparsevllm/operators/moe.py +++ b/src/sparsevllm/operators/moe.py @@ -6,6 +6,7 @@ import torch import sparsevllm.platforms as platforms +from sparsevllm.kernels.moe import MoeAlignment from sparsevllm.operators.registry import ( OpRegistry, OpResolver, @@ -176,6 +177,33 @@ def run( MOE_REGISTRY: OpRegistry[MoeOpSpec, MoeProvider] = OpRegistry("routed MoE") +def _sgl_moe_align_block_size( + topk_ids: torch.Tensor, + *, + block_size: int, + num_experts: int, + local_expert_start: int, + local_expert_end: int, +) -> MoeAlignment: + from sparsevllm.kernels.external.sgl.moe import sgl_moe_align_block_size + from sparsevllm.kernels.triton.moe import localize_expert_ids + + num_local_experts = int(local_expert_end) - int(local_expert_start) + has_remote_experts = num_local_experts != int(num_experts) + if has_remote_experts: + topk_ids = localize_expert_ids( + topk_ids, + local_expert_start=local_expert_start, + local_expert_end=local_expert_end, + remote_expert_id=num_local_experts, + ) + return sgl_moe_align_block_size( + topk_ids, + block_size=block_size, + num_experts=num_local_experts + int(has_remote_experts), + ) + + @MOE_REGISTRY.register class SglAlignedTritonGlmMoeProvider(MoeProvider): name = "sgl_aligned_triton_glm" @@ -211,7 +239,7 @@ def supports(cls, spec: MoeOpSpec, caps: DeviceCaps) -> SupportResult: return SupportResult.no("requires BF16 activations") if spec.weight_dtype != torch.bfloat16 or spec.block_shape is not None: return SupportResult.no("requires unquantized BF16 expert weights") - from sparsevllm.operators.sgl_moe import sgl_moe_alignment_support + from sparsevllm.kernels.external.sgl.moe import sgl_moe_alignment_support supported, reason = sgl_moe_alignment_support() return SupportResult.yes() if supported else SupportResult.no(reason) @@ -233,8 +261,7 @@ def run( del ep_rank if w13_scale_inv is not None or w2_scale_inv is not None: raise RuntimeError("SGL-aligned BF16 MoE does not accept expert scales.") - from sparsevllm.operators.sgl_moe import sgl_moe_align_block_size - from sparsevllm.triton_kernel.moe import fused_moe + from sparsevllm.kernels.triton.moe import fused_moe alignment_impl = None num_tokens = int(hidden_states.shape[0]) @@ -245,7 +272,7 @@ def run( or num_tokens * int(spec.top_k) * 4 > int(spec.num_experts) ) ): - alignment_impl = sgl_moe_align_block_size + alignment_impl = _sgl_moe_align_block_size return fused_moe( hidden_states, w13_weight, @@ -343,7 +370,7 @@ def run( del ep_rank if w13_scale_inv is None or w2_scale_inv is None: raise RuntimeError("MiniMax M2.7 fused MoE requires expert scales.") - from sparsevllm.triton_kernel.minimax_m2_moe import ( + from sparsevllm.kernels.triton.minimax_m2_moe import ( fused_minimax_m2_moe_fp8, ) @@ -502,7 +529,7 @@ def run( del ep_rank if w13_scale_inv is not None or w2_scale_inv is not None: raise RuntimeError("Fused Hopper BF16 MoE does not accept expert scales.") - from sparsevllm.triton_kernel.moe import fused_moe_gate_up_swiglu + from sparsevllm.kernels.triton.moe import fused_moe_gate_up_swiglu return fused_moe_gate_up_swiglu( hidden_states, @@ -577,7 +604,7 @@ def run( if spec.weight_dtype == torch.float8_e4m3fn: if w13_scale_inv is None or w2_scale_inv is None: raise RuntimeError("Triton FP8 MoE requires expert scales.") - from sparsevllm.triton_kernel.moe import fused_moe_fp8 + from sparsevllm.kernels.triton.moe import fused_moe_fp8 return fused_moe_fp8( hidden_states, @@ -591,7 +618,7 @@ def run( local_expert_start=local_expert_start, gate_up_order=self.gate_up_order, ) - from sparsevllm.triton_kernel.moe import fused_moe + from sparsevllm.kernels.triton.moe import fused_moe return fused_moe( hidden_states, @@ -686,7 +713,7 @@ def run( ) if w13_scale_inv is None or w2_scale_inv is None: raise RuntimeError("Qwen3.6 hybrid FP8 MoE requires expert scales.") - from sparsevllm.triton_kernel.moe import fused_moe_fp8 + from sparsevllm.kernels.triton.moe import fused_moe_fp8 return fused_moe_fp8( hidden_states, diff --git a/src/sparsevllm/operators/moe_router.py b/src/sparsevllm/operators/moe_router.py index 148106ec..19e4d457 100644 --- a/src/sparsevllm/operators/moe_router.py +++ b/src/sparsevllm/operators/moe_router.py @@ -98,7 +98,7 @@ def run( ) -> tuple[torch.Tensor, torch.Tensor]: if correction_bias is not None or routed_scaling_factor != 1.0: raise ValueError("Softmax routing does not accept bias or route scaling.") - from sparsevllm.triton_kernel.moe_topk import topk_softmax + from sparsevllm.kernels.triton.moe_topk import topk_softmax return topk_softmax( router_logits, @@ -134,7 +134,7 @@ def run( ) -> tuple[torch.Tensor, torch.Tensor]: if correction_bias is None: raise ValueError("Biased-sigmoid routing requires correction_bias.") - from sparsevllm.triton_kernel.moe_biased_sigmoid import ( + from sparsevllm.kernels.triton.moe_biased_sigmoid import ( fused_topk_biased_sigmoid, ) diff --git a/src/sparsevllm/operators/prefill_attention.py b/src/sparsevllm/operators/prefill_attention.py index 921dffb2..76cea45a 100644 --- a/src/sparsevllm/operators/prefill_attention.py +++ b/src/sparsevllm/operators/prefill_attention.py @@ -369,7 +369,7 @@ def run( del spec, layer_idx payload, meta = _view_parts(view) _validate_token_page_table(meta) - from sparsevllm.triton_kernel.context_flashattention_nopad import ( + from sparsevllm.kernels.triton.context_flashattention_nopad import ( context_attention_fwd, ) diff --git a/tests/test_deltakv_less_memory_kernel.py b/tests/test_deltakv_less_memory_kernel.py index 0cba6ea9..0cf412b3 100644 --- a/tests/test_deltakv_less_memory_kernel.py +++ b/tests/test_deltakv_less_memory_kernel.py @@ -4,7 +4,7 @@ import torch -from sparsevllm.triton_kernel.deltakv_kernels import ( +from sparsevllm.kernels.triton.deltakv_kernels import ( _validate_full_layer_kivi_decode_maps, deltakv_less_memory_reconstruct_writeback_quantized, deltakv_l2_topk_blockwise, @@ -18,11 +18,11 @@ full_layer_kivi_flash_decode_stage1_token_group_map, full_layer_kivi_flash_decode_stage1_token_map, ) -from sparsevllm.triton_kernel.gqa_flash_decoding_stage1 import flash_decode_stage1 as gqa_flash_decode_stage1 -from sparsevllm.triton_kernel.gqa_flash_decoding_stage1 import ( +from sparsevllm.kernels.triton.gqa_flash_decoding_stage1 import flash_decode_stage1 as gqa_flash_decode_stage1 +from sparsevllm.kernels.triton.gqa_flash_decoding_stage1 import ( flash_decode_stage1_with_score as gqa_flash_decode_stage1_with_score, ) -from sparsevllm.triton_kernel.quant import ( +from sparsevllm.kernels.triton.quant import ( triton_dequantize_2d_int4_grouped, triton_quantize_and_pack_2d_int4_grouped, triton_quantize_and_pack_along_last_dim, @@ -95,7 +95,7 @@ def test_residual_2d_int4_dequant_matches_reference_for_float32(self): self.assertTrue(torch.allclose(got, ref, atol=1e-6, rtol=1e-6)) def test_store_kvcache_chunks_large_launches(self): - from sparsevllm.triton_kernel.store_kvcache import store_kvcache + from sparsevllm.kernels.triton.store_kvcache import store_kvcache old_value = os.environ.get("SPARSEVLLM_STORE_KVCACHE_CHUNK_TOKENS") os.environ["SPARSEVLLM_STORE_KVCACHE_CHUNK_TOKENS"] = "3" diff --git a/tests/test_minimax_m2_attention_graph.py b/tests/test_minimax_m2_attention_graph.py index b3b5c0a2..b41cb20e 100644 --- a/tests/test_minimax_m2_attention_graph.py +++ b/tests/test_minimax_m2_attention_graph.py @@ -1,9 +1,9 @@ import pytest import torch -from sparsevllm.triton_kernel.flash_decoding_stage2 import flash_decode_stage2 -from sparsevllm.triton_kernel.gqa_flash_decoding_stage1 import flash_decode_stage1 -from sparsevllm.triton_kernel.store_kvcache import store_kvcache +from sparsevllm.kernels.triton.flash_decoding_stage2 import flash_decode_stage2 +from sparsevllm.kernels.triton.gqa_flash_decoding_stage1 import flash_decode_stage1 +from sparsevllm.kernels.triton.store_kvcache import store_kvcache @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") diff --git a/tests/test_minimax_m2_router.py b/tests/test_minimax_m2_router.py index 8abe58a3..c3d281a9 100644 --- a/tests/test_minimax_m2_router.py +++ b/tests/test_minimax_m2_router.py @@ -4,11 +4,11 @@ import torch import torch.nn.functional as F -from sparsevllm.triton_kernel.minimax_m2_router import ( +from sparsevllm.kernels.triton.minimax_m2_router import ( minimax_m2_router, topk_biased_sigmoid, ) -from sparsevllm.triton_kernel.moe_biased_sigmoid import ( +from sparsevllm.kernels.triton.moe_biased_sigmoid import ( topk_biased_sigmoid as generic_topk_biased_sigmoid, ) diff --git a/tests/test_mla_attention_operator.py b/tests/test_mla_attention_operator.py index 56f86fa7..42996b89 100644 --- a/tests/test_mla_attention_operator.py +++ b/tests/test_mla_attention_operator.py @@ -21,7 +21,7 @@ ) from sparsevllm.operators.registry import OpResolver from sparsevllm.platforms import DeviceCaps, PlatformEnum -from sparsevllm.triton_kernel.mla import ( +from sparsevllm.kernels.triton.mla import ( GLM_MLA_MAX_WORKSPACE_CONFIG, MlaDecodeWorkspace, ) diff --git a/tests/test_mla_kernels.py b/tests/test_mla_kernels.py index 6df057ce..5cfd01f8 100644 --- a/tests/test_mla_kernels.py +++ b/tests/test_mla_kernels.py @@ -7,7 +7,7 @@ import pytest import torch -from sparsevllm.triton_kernel.mla import ( +from sparsevllm.kernels.triton.mla import ( DEFAULT_GLM_MLA_DECODE_CONFIG, GLM_MLA_SOFTMAX_SCALE, MlaDecodeWorkspace, @@ -159,7 +159,8 @@ def test_vendor_python_files_do_not_import_lightllm() -> None: Path(__file__).parents[1] / "src" / "sparsevllm" - / "triton_kernel" + / "kernels" + / "triton" / "mla" ) for path in kernel_dir.glob("*.py"): diff --git a/tests/test_moe_config.py b/tests/test_moe_config.py index 868f59c2..b2776447 100644 --- a/tests/test_moe_config.py +++ b/tests/test_moe_config.py @@ -1,7 +1,7 @@ import pytest import torch -from sparsevllm.triton_kernel.moe_config import ( +from sparsevllm.kernels.triton.moe_config import ( resolve_fp8_routed_gemm_config, resolve_moe_gemm_config, token_bucket, diff --git a/tests/test_operator_providers.py b/tests/test_operator_providers.py index 9b06153d..9a50d8d2 100644 --- a/tests/test_operator_providers.py +++ b/tests/test_operator_providers.py @@ -582,7 +582,7 @@ def test_glm_tp2_moe_uses_sgl_aligned_provider(): ) with patch( - "sparsevllm.operators.sgl_moe.sgl_moe_alignment_support", + "sparsevllm.kernels.external.sgl.moe.sgl_moe_alignment_support", return_value=(True, "available"), ): resolved = OpResolver(MOE_REGISTRY).resolve( @@ -613,7 +613,7 @@ def test_glm_tp2_fused_shared_decode_uses_sgl_aligned_provider(): ) with patch( - "sparsevllm.operators.sgl_moe.sgl_moe_alignment_support", + "sparsevllm.kernels.external.sgl.moe.sgl_moe_alignment_support", return_value=(True, "available"), ): resolved = OpResolver(MOE_REGISTRY).resolve( @@ -644,7 +644,7 @@ def test_glm_tp2_ep2_moe_uses_sgl_aligned_provider(): ) with patch( - "sparsevllm.operators.sgl_moe.sgl_moe_alignment_support", + "sparsevllm.kernels.external.sgl.moe.sgl_moe_alignment_support", return_value=(True, "supported"), ): resolved = OpResolver(MOE_REGISTRY).resolve( @@ -686,7 +686,7 @@ def test_glm_tp2_ep2_sgl_alignment_is_bounded( weights = torch.empty(0) provider = SglAlignedTritonGlmMoeProvider() - with patch("sparsevllm.triton_kernel.moe.fused_moe") as fused_moe: + with patch("sparsevllm.kernels.triton.moe.fused_moe") as fused_moe: provider.run( spec, hidden_states, @@ -1044,7 +1044,7 @@ def test_qwen36_hybrid_moe_dispatches_by_token_bucket(): with patch.dict( sys.modules, { - "sparsevllm.triton_kernel.moe": SimpleNamespace( + "sparsevllm.kernels.triton.moe": SimpleNamespace( fused_moe_fp8=triton_call ) }, @@ -1105,7 +1105,7 @@ def test_qwen36_hybrid_moe_uses_larger_triton_bucket_on_single_gpu(): with patch.dict( sys.modules, { - "sparsevllm.triton_kernel.moe": SimpleNamespace( + "sparsevllm.kernels.triton.moe": SimpleNamespace( fused_moe_fp8=triton_call ) }, diff --git a/tests/test_prefill_score_kernel.py b/tests/test_prefill_score_kernel.py index 9de76dd7..a5cc3372 100644 --- a/tests/test_prefill_score_kernel.py +++ b/tests/test_prefill_score_kernel.py @@ -61,7 +61,7 @@ def _prefill_score_baseline( @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for prefill score Triton tests.") class PrefillScoreKernelTest(unittest.TestCase): def test_prefill_score_matches_torch_for_query_range(self): - from sparsevllm.triton_kernel.prefill_score import prefill_score_fwd + from sparsevllm.kernels.triton.prefill_score import prefill_score_fwd torch.manual_seed(7) device = "cuda" @@ -121,7 +121,7 @@ def test_prefill_score_matches_torch_for_query_range(self): torch.testing.assert_close(attn_score, expected, rtol=2e-2, atol=2e-2) def test_prefill_score_handles_offset_query_window(self): - from sparsevllm.triton_kernel.prefill_score import prefill_score_fwd + from sparsevllm.kernels.triton.prefill_score import prefill_score_fwd torch.manual_seed(11) device = "cuda" @@ -179,7 +179,7 @@ def test_prefill_score_handles_offset_query_window(self): torch.testing.assert_close(acc, expected, rtol=2e-2, atol=2e-2) def test_prefill_score_matches_torch_for_gqa_seven_heads(self): - from sparsevllm.triton_kernel.prefill_score import prefill_score_fwd + from sparsevllm.kernels.triton.prefill_score import prefill_score_fwd torch.manual_seed(17) device = "cuda" diff --git a/tests/test_qwen35_causal_conv1d.py b/tests/test_qwen35_causal_conv1d.py index 3a6395b1..27b64925 100644 --- a/tests/test_qwen35_causal_conv1d.py +++ b/tests/test_qwen35_causal_conv1d.py @@ -2,7 +2,7 @@ import torch import torch.nn.functional as F -from sparsevllm.triton_kernel.qwen3_5.causal_conv1d import causal_conv1d_fn +from sparsevllm.kernels.triton.qwen3_5.causal_conv1d import causal_conv1d_fn pytestmark = pytest.mark.skipif( diff --git a/tests/test_qwen35_hd256_decode_routing.py b/tests/test_qwen35_hd256_decode_routing.py index 09daeb27..99fb48bb 100644 --- a/tests/test_qwen35_hd256_decode_routing.py +++ b/tests/test_qwen35_hd256_decode_routing.py @@ -9,8 +9,8 @@ ExplicitKVPayload, ) from sparsevllm.layers.attention_backend import TritonAttentionBackend -from sparsevllm.triton_kernel.flash_decoding_stage2 import flash_decode_stage2 -from sparsevllm.triton_kernel.gqa_flash_decoding_stage1 import ( +from sparsevllm.kernels.triton.flash_decoding_stage2 import flash_decode_stage2 +from sparsevllm.kernels.triton.gqa_flash_decoding_stage1 import ( flash_decode_stage1, flash_decode_stage1_with_score, ) @@ -179,7 +179,7 @@ def test_grouped_gqa_wrappers_accept_head_dim_256(self): attn_score = torch.empty(1, 24, 3) with patch( - "sparsevllm.triton_kernel.gqa_flash_decoding_stage1._fwd_kernel_flash_decode_stage1" + "sparsevllm.kernels.triton.gqa_flash_decoding_stage1._fwd_kernel_flash_decode_stage1" ) as kernel: flash_decode_stage1( q, @@ -196,7 +196,7 @@ def test_grouped_gqa_wrappers_accept_head_dim_256(self): kernel.__getitem__.return_value.assert_called_once() with patch( - "sparsevllm.triton_kernel.gqa_flash_decoding_stage1._fwd_kernel_flash_decode_stage1_with_score" + "sparsevllm.kernels.triton.gqa_flash_decoding_stage1._fwd_kernel_flash_decode_stage1_with_score" ) as kernel: flash_decode_stage1_with_score( q, @@ -261,7 +261,7 @@ def test_unified_stage2_forwards_noncontiguous_strides_and_selects_warps(self): b_seqlen = torch.tensor([257], dtype=torch.int32) output = torch.empty(1, 24, head_dim * 2)[..., ::2] with patch( - "sparsevllm.triton_kernel.flash_decoding_stage2._fwd_kernel_flash_decode_stage2" + "sparsevllm.kernels.triton.flash_decoding_stage2._fwd_kernel_flash_decode_stage2" ) as kernel: flash_decode_stage2(mid_out, mid_lse, b_seqlen, output, 256) diff --git a/tests/test_qwen3_moe.py b/tests/test_qwen3_moe.py index 2c836fba..69d65235 100644 --- a/tests/test_qwen3_moe.py +++ b/tests/test_qwen3_moe.py @@ -269,7 +269,7 @@ def save_rope_kv_if_needed(self, _layer_idx, _key, _value): def test_moe_block_uses_triton_kernels(): - from sparsevllm.triton_kernel.moe_topk import topk_softmax + from sparsevllm.kernels.triton.moe_topk import topk_softmax config = _config() context = _ep_context(0, 1) diff --git a/tests/test_sgl_fa3.py b/tests/test_sgl_fa3.py index cbb7fb28..2f211877 100644 --- a/tests/test_sgl_fa3.py +++ b/tests/test_sgl_fa3.py @@ -5,7 +5,7 @@ import pytest import torch -from sparsevllm.operators.sgl_fa3 import ( +from sparsevllm.kernels.external.sgl.fa3 import ( SglFa3DecodeKernel, sgl_fa3_support, ) diff --git a/tests/test_sgl_moe.py b/tests/test_sgl_moe.py index f60c29c9..932b891b 100644 --- a/tests/test_sgl_moe.py +++ b/tests/test_sgl_moe.py @@ -5,12 +5,12 @@ import pytest import torch -from sparsevllm.operators.sgl_moe import ( - sgl_moe_alignment_support, +from sparsevllm.kernels.external.sgl.moe import ( sgl_moe_align_block_size, + sgl_moe_alignment_support, ) -from sparsevllm.triton_kernel.moe import moe_align_block_size -from sparsevllm.triton_kernel.moe import fused_moe +from sparsevllm.kernels.triton.moe import fused_moe, moe_align_block_size +from sparsevllm.operators.moe import _sgl_moe_align_block_size def test_sgl_moe_support_rejects_missing_package() -> None: @@ -66,8 +66,6 @@ def test_sgl_moe_alignment_matches_grouped_assignments() -> None: topk_ids, block_size=16, num_experts=64, - local_expert_start=0, - local_expert_end=64, ) torch.cuda.synchronize() count = int(reference_count.item()) @@ -117,7 +115,7 @@ def test_sgl_moe_alignment_matches_ep_shard( local_expert_start=local_start, local_expert_end=local_end, ) - actual = sgl_moe_align_block_size( + actual = _sgl_moe_align_block_size( topk_ids, block_size=16, num_experts=64, @@ -242,7 +240,7 @@ def test_sgl_ep_alignment_preserves_full_fused_moe_output( w2_weight, topk_ids, topk_weights, - alignment_impl=sgl_moe_align_block_size, + alignment_impl=_sgl_moe_align_block_size, **kwargs, ) torch.cuda.synchronize() @@ -258,7 +256,7 @@ def test_sgl_ep_alignment_preserves_full_fused_moe_output( w2_weight, topk_ids, topk_weights, - alignment_impl=sgl_moe_align_block_size, + alignment_impl=_sgl_moe_align_block_size, **kwargs, ) hidden_states.copy_(torch.randn_like(hidden_states)) diff --git a/tests/test_tilelang_mla_kernel.py b/tests/test_tilelang_mla_kernel.py index 5055be57..6cec7701 100644 --- a/tests/test_tilelang_mla_kernel.py +++ b/tests/test_tilelang_mla_kernel.py @@ -3,12 +3,11 @@ import pytest import torch -from sparsevllm.operators.tilelang_mla import ( +from sparsevllm.kernels.tilelang.mla.runtime import ( TileMlaDecodeKernel, TileMlaLaunchConfig, ) - CUDA_REQUIRED = pytest.mark.skipif( not torch.cuda.is_available(), reason="CUDA is required for the TileLang MLA kernel test", diff --git a/tests/test_tilelang_mla_operator.py b/tests/test_tilelang_mla_operator.py index 1d8db2e7..8dbd059c 100644 --- a/tests/test_tilelang_mla_operator.py +++ b/tests/test_tilelang_mla_operator.py @@ -11,6 +11,12 @@ DecodeComputeView, MlaLatentPayload, ) +from sparsevllm.kernels.tilelang.mla.runtime import ( + TileMlaDecodeKernel, + select_tile_mla_config, + tilelang_mla_support, +) +from sparsevllm.kernels.triton.mla import MlaDecodeWorkspace from sparsevllm.operators.mla_attention import ( MLA_ATTENTION_REGISTRY, MlaAttentionOpSpec, @@ -18,13 +24,7 @@ MlaTileLangScoreProvider, ) from sparsevllm.operators.registry import OpResolver -from sparsevllm.operators.tilelang_mla import ( - TileMlaDecodeKernel, - select_tile_mla_config, - tilelang_mla_support, -) from sparsevllm.platforms import DeviceCaps, PlatformEnum -from sparsevllm.triton_kernel.mla import MlaDecodeWorkspace def _spec(*, tp_size: int = 2) -> MlaAttentionOpSpec: diff --git a/tests/test_triton_fp8_operators.py b/tests/test_triton_fp8_operators.py index fafc7e14..32896367 100644 --- a/tests/test_triton_fp8_operators.py +++ b/tests/test_triton_fp8_operators.py @@ -4,9 +4,9 @@ from sparsevllm.operators.fp8_linear import resolve_fp8_linear_provider from sparsevllm.quantization.fp8 import fp8_blockwise_linear_reference -from sparsevllm.triton_kernel.fp8_blockwise import fp8_blockwise_matmul -from sparsevllm.triton_kernel.moe import fused_moe_fp8 -from sparsevllm.triton_kernel.minimax_m2_moe import fused_minimax_m2_moe_fp8 +from sparsevllm.kernels.triton.fp8_blockwise import fp8_blockwise_matmul +from sparsevllm.kernels.triton.moe import fused_moe_fp8 +from sparsevllm.kernels.triton.minimax_m2_moe import fused_minimax_m2_moe_fp8 pytestmark = pytest.mark.skipif( diff --git a/tests/test_triton_moe.py b/tests/test_triton_moe.py index fdd6cd7f..840b1c40 100644 --- a/tests/test_triton_moe.py +++ b/tests/test_triton_moe.py @@ -5,16 +5,16 @@ import torch.nn.functional as F from sparsevllm.operators.gated_shared_add import gated_shared_add -from sparsevllm.triton_kernel.gate_up_swiglu import h20_gate_up_swiglu -from sparsevllm.triton_kernel.moe import ( +from sparsevllm.kernels.triton.gate_up_swiglu import h20_gate_up_swiglu +from sparsevllm.kernels.triton.moe import ( _prepare_expert_assignment, append_shared_expert_route, fused_moe, fused_moe_gate_up_swiglu, moe_align_block_size, ) -from sparsevllm.triton_kernel.moe_topk import topk_softmax -from sparsevllm.triton_kernel.silu_and_mul import _resolve_silu_launch_config +from sparsevllm.kernels.triton.moe_topk import topk_softmax +from sparsevllm.kernels.triton.silu_and_mul import _resolve_silu_launch_config def test_silu_launch_config_uses_decode_tile_only_for_small_rows(): From 664bc6986cf8df6554ac2ea04edb9c4c07873bf1 Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Wed, 12 Aug 2026 19:16:58 +0800 Subject: [PATCH 35/48] refactor: simplify cuda graph runner --- src/sparsevllm/engine/decode_cuda_graph.py | 18 +++++------------- src/sparsevllm/engine/model_runner.py | 2 -- tests/test_glm_cuda_graph.py | 6 ------ 3 files changed, 5 insertions(+), 21 deletions(-) diff --git a/src/sparsevllm/engine/decode_cuda_graph.py b/src/sparsevllm/engine/decode_cuda_graph.py index a950bbe9..6ac0e70d 100644 --- a/src/sparsevllm/engine/decode_cuda_graph.py +++ b/src/sparsevllm/engine/decode_cuda_graph.py @@ -88,8 +88,6 @@ def __init__( run_model: Callable[[torch.Tensor, torch.Tensor, bool], torch.Tensor], is_long_text_batch: Callable[[list[Sequence], bool], bool], method: str, - rank: int, - parallel_context, capture_sizes: list[int], context_sizes: list[int] | tuple[int, ...] | str | int | None = None, graph_pool=None, @@ -101,8 +99,6 @@ def __init__( self.run_model = run_model self.is_long_text_batch = is_long_text_batch self.method = str(method or "") - self.rank = int(rank) - self.parallel_context = parallel_context self.platform = platforms.current_platform self.capture_sizes = sorted(set(int(size) for size in capture_sizes)) if not self.capture_sizes or any(size <= 0 for size in self.capture_sizes): @@ -478,7 +474,7 @@ def _capture( if sparse_keepalive is not None: keepalive.extend(sparse_keepalive()) state.keepalive = keepalive - self.capture_count = int(getattr(self, "capture_count", 0)) + 1 + self.capture_count += 1 return state def run( @@ -503,9 +499,7 @@ def run( real_batch_size = len(seqs) force_eager = getattr(self.cache_manager, "decode_cuda_graph_force_eager", None) if force_eager is not None and force_eager(): - self.force_eager_count = int( - getattr(self, "force_eager_count", 0) - ) + 1 + self.force_eager_count += 1 return self.run_eager_static(seqs), None graph_batch_size = self._select_graph_batch_size(real_batch_size) @@ -528,7 +522,7 @@ def run( self._restore_sparse_state_refs(state) with profiler.record("decode_cuda_graph_replay_after_capture"): state.graph.replay() - self.replay_count = int(getattr(self, "replay_count", 0)) + 1 + self.replay_count += 1 logits = state.logits[:real_batch_size] if state.logits is not None else None token_ids = state.token_ids[:real_batch_size] if state.token_ids is not None else None return logits, token_ids @@ -536,7 +530,7 @@ def run( self._restore_sparse_state_refs(state) with profiler.record("decode_cuda_graph_replay"): state.graph.replay() - self.replay_count = int(getattr(self, "replay_count", 0)) + 1 + self.replay_count += 1 logits = state.logits[:real_batch_size] if state.logits is not None else None token_ids = state.token_ids[:real_batch_size] if state.token_ids is not None else None return logits, token_ids @@ -545,9 +539,7 @@ def run_eager_static(self, seqs: list[Sequence]) -> torch.Tensor | None: """Run decode eagerly through the same static-compatible path used by graphs.""" if not seqs: raise ValueError("static decode requires a non-empty decode batch.") - self.eager_static_count = int( - getattr(self, "eager_static_count", 0) - ) + 1 + self.eager_static_count += 1 real_batch_size = len(seqs) graph_batch_size = self._select_graph_batch_size(real_batch_size) diff --git a/src/sparsevllm/engine/model_runner.py b/src/sparsevllm/engine/model_runner.py index 9293c9fc..c688334f 100644 --- a/src/sparsevllm/engine/model_runner.py +++ b/src/sparsevllm/engine/model_runner.py @@ -297,8 +297,6 @@ def __init__( run_model=self.run_model, is_long_text_batch=self._is_long_text_batch, method=self.config.vllm_sparse_method, - rank=self.rank, - parallel_context=self.parallel_context, capture_sizes=decode_static_capture_sizes, context_sizes=decode_static_context_sizes, graph_pool=self.cuda_graph_pool, diff --git a/tests/test_glm_cuda_graph.py b/tests/test_glm_cuda_graph.py index 6dda45da..7fdbe739 100644 --- a/tests/test_glm_cuda_graph.py +++ b/tests/test_glm_cuda_graph.py @@ -248,8 +248,6 @@ def run_model( run_model=run_model, is_long_text_batch=lambda seqs, is_prefill: False, method="", - rank=0, - parallel_context=parallel_context, capture_sizes=[1], context_sizes=[128], ) @@ -558,8 +556,6 @@ def run_model( run_model=run_model, is_long_text_batch=lambda seqs, is_prefill: False, method="", - rank=0, - parallel_context=parallel_context, capture_sizes=[1], context_sizes=[128], ) @@ -1118,8 +1114,6 @@ def run_model( run_model=run_model, is_long_text_batch=lambda seqs, is_prefill: True, method=method, - rank=0, - parallel_context=parallel_context, capture_sizes=[1], context_sizes=[16], ) From 6798c5b2030d752685e5f2813c50f5fe429990d0 Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Wed, 12 Aug 2026 19:37:10 +0800 Subject: [PATCH 36/48] refactor: share packed moe experts --- src/sparsevllm/layers/packed_moe.py | 11 +- src/sparsevllm/models/minimax_m2.py | 191 +++------------------------- tests/test_minimax_m2_model.py | 2 +- tests/test_packed_moe.py | 4 +- 4 files changed, 27 insertions(+), 181 deletions(-) diff --git a/src/sparsevllm/layers/packed_moe.py b/src/sparsevllm/layers/packed_moe.py index 0356c17d..9a4d0f2e 100644 --- a/src/sparsevllm/layers/packed_moe.py +++ b/src/sparsevllm/layers/packed_moe.py @@ -26,6 +26,7 @@ class PackedMoeExperts(PackedExpertWeightLoader, nn.Module): "up_proj": "up", "down_proj": "down", } + checkpoint_scale_dtype = torch.bfloat16 def __init__( self, @@ -64,6 +65,7 @@ def __init__( self.global_intermediate_size // self.tp_size ) self.fp8_enabled = bool(fp8_enabled) + scale_dtype = (scale_dtype or torch.float32) if self.fp8_enabled else None if self.fp8_enabled and ( self.hidden_size % 128 or self.global_intermediate_size % 128 ): @@ -146,7 +148,7 @@ def __init__( self.num_local_experts, 2 * self.intermediate_size // 128, self.hidden_size // 128, - dtype=torch.float32, + dtype=scale_dtype, ), ) self.register_buffer( @@ -155,7 +157,7 @@ def __init__( self.num_local_experts, self.hidden_size // 128, self.intermediate_size // 128, - dtype=torch.float32, + dtype=scale_dtype, ), ) else: @@ -200,9 +202,10 @@ def load_expert_weight( f"{self.model_label} expert weight must be FP8 E4M3, " f"got {loaded_weight.dtype}." ) - if loaded_scale.dtype != torch.bfloat16: + if loaded_scale.dtype != self.checkpoint_scale_dtype: raise TypeError( - f"{self.model_label} expert weight_scale_inv must be BF16, " + f"{self.model_label} expert weight_scale_inv must use " + f"{self.checkpoint_scale_dtype}, " f"got {loaded_scale.dtype}." ) elif loaded_scale is not None: diff --git a/src/sparsevllm/models/minimax_m2.py b/src/sparsevllm/models/minimax_m2.py index 1efcd448..7d9a8067 100644 --- a/src/sparsevllm/models/minimax_m2.py +++ b/src/sparsevllm/models/minimax_m2.py @@ -15,19 +15,15 @@ ) from sparsevllm.layers.attention import Attention from sparsevllm.layers.embed_head import ParallelLMHead -from sparsevllm.layers.expert_weights import PackedExpertWeightLoader from sparsevllm.layers.layernorm import ColumnParallelRMSNorm, RMSNorm from sparsevllm.layers.linear import QKVParallelLinear, RowParallelLinear +from sparsevllm.layers.packed_moe import PackedMoeExperts from sparsevllm.layers.rotary_embedding import ( apply_partial_rotary_emb, get_rope, ) from sparsevllm.models.qwen3 import Qwen3ModelBase -from sparsevllm.operators.moe import ( - MoeOpSpec, - model_activation_dtype, - resolve_moe_provider, -) +from sparsevllm.operators.moe import model_activation_dtype, resolve_moe_provider from sparsevllm.operators.all_reduce import ( AllReduceOpSpec, PreparedAllReduceOp, @@ -44,7 +40,6 @@ prepare_prefill_attention_op, ) from sparsevllm.platforms import device_runtime -from sparsevllm.quantization.fp8_tp import Fp8ExpertTpShard from sparsevllm.utils.context import get_context from sparsevllm.utils.log import logger from sparsevllm.utils.weight_target import WeightTarget @@ -218,8 +213,9 @@ def forward( return router_logits, topk_weights, topk_ids -class MiniMaxM2PackedExperts(PackedExpertWeightLoader, nn.Module): +class MiniMaxM2PackedExperts(PackedMoeExperts): checkpoint_projection_map = {"w1": "gate", "w2": "down", "w3": "up"} + checkpoint_scale_dtype = torch.float32 def __init__( self, @@ -227,184 +223,29 @@ def __init__( *, cuda_graph: bool | None = None, ) -> None: - super().__init__() - parallel_context = get_parallel_context() - self.tp_rank = int(parallel_context.moe_tp_rank) - self.tp_size = int(parallel_context.moe_tp_size) - self.ep_rank = int(parallel_context.ep_rank) - self.ep_size = int(parallel_context.ep_size) - self.num_experts = int(config.num_local_experts) - self.hidden_size = int(config.hidden_size) - self.global_intermediate_size = int(config.intermediate_size) - self.fp8_tp_shard = Fp8ExpertTpShard( - self.global_intermediate_size, - self.tp_rank, - self.tp_size, - ) - self.checkpoint_tp_shard = self.fp8_tp_shard - self.logical_intermediate_size = self.fp8_tp_shard.logical_size - self.intermediate_size = self.fp8_tp_shard.physical_size - if self.num_experts % self.ep_size: - raise ValueError( - f"MiniMax experts={self.num_experts} must be divisible by EP={self.ep_size}." - ) - if self.hidden_size % 128 or self.global_intermediate_size % 128: + block_shape = tuple(config.quantization_config.weight_block_size) + if block_shape != (128, 128): raise ValueError( - "MiniMax packed FP8 experts require hidden/intermediate dimensions " - f"aligned to 128, got {self.hidden_size}/{self.global_intermediate_size}." + "MiniMax packed FP8 experts require weight_block_size=(128, 128), " + f"got {block_shape}." ) - self.num_local_experts = self.num_experts // self.ep_size - self.local_expert_start = self.ep_rank * self.num_local_experts - self.local_expert_end = self.local_expert_start + self.num_local_experts - self.op_spec = MoeOpSpec( - num_experts=self.num_experts, - num_local_experts=self.num_local_experts, - hidden_size=self.hidden_size, - intermediate_size=self.intermediate_size, + super().__init__( + num_experts=int(config.num_local_experts), + hidden_size=int(config.hidden_size), + intermediate_size=int(config.intermediate_size), top_k=int(config.num_experts_per_tok), activation_dtype=model_activation_dtype(config), - weight_dtype=torch.float8_e4m3fn, - block_shape=tuple(config.quantization_config.weight_block_size), - ep_size=self.ep_size, + fp8_enabled=True, cuda_graph=( bool(getattr(config, "decode_cuda_graph", False)) if cuda_graph is None else bool(cuda_graph) ), - tp_size=self.tp_size, routing_method="biased_sigmoid", scale_dtype=torch.float32, - ) - self.provider = resolve_moe_provider(self.op_spec) - self.w13_weight = nn.Parameter( - torch.empty( - self.num_local_experts, - 2 * self.intermediate_size, - self.hidden_size, - dtype=torch.float8_e4m3fn, - ), - requires_grad=False, - ) - self.w2_weight = nn.Parameter( - torch.empty( - self.num_local_experts, - self.hidden_size, - self.intermediate_size, - dtype=torch.float8_e4m3fn, - ), - requires_grad=False, - ) - self.register_buffer( - "w13_scale_inv", - torch.empty( - self.num_local_experts, - 2 * self.intermediate_size // 128, - self.hidden_size // 128, - dtype=torch.float32, - ), - ) - self.register_buffer( - "w2_scale_inv", - torch.empty( - self.num_local_experts, - self.hidden_size // 128, - self.intermediate_size // 128, - dtype=torch.float32, - ), - ) - self._loaded_expert_shards: set[tuple[int, str]] = set() - - def is_local_expert(self, global_expert_id: int) -> bool: - return self.local_expert_start <= int(global_expert_id) < self.local_expert_end - - def load_expert_weight( - self, - global_expert_id: int, - projection: str, - loaded_weight: torch.Tensor, - loaded_scale: torch.Tensor | None, - ) -> None: - global_expert_id = int(global_expert_id) - if not self.is_local_expert(global_expert_id): - raise ValueError( - f"Expert {global_expert_id} is outside local range " - f"[{self.local_expert_start}, {self.local_expert_end})." - ) - if projection not in {"w1", "w2", "w3"}: - raise ValueError(f"Unsupported MiniMax expert projection {projection!r}.") - load_key = (global_expert_id, projection) - if load_key in self._loaded_expert_shards: - raise ValueError( - f"Duplicate MiniMax expert weight for expert={global_expert_id}, " - f"projection={projection}." - ) - if loaded_scale is None: - raise ValueError( - f"Missing FP8 weight_scale_inv for MiniMax expert={global_expert_id}, " - f"projection={projection}." - ) - if loaded_weight.dtype != torch.float8_e4m3fn: - raise TypeError( - f"MiniMax expert weight must be FP8 E4M3, got {loaded_weight.dtype}." - ) - if loaded_scale.dtype != torch.float32: - raise TypeError( - "MiniMax expert weight_scale_inv must be FP32, " - f"got {loaded_scale.dtype}." - ) - - loaded_weight, loaded_scale = self.fp8_tp_shard.prepare_projection( - loaded_weight, - loaded_scale, - hidden_size=self.hidden_size, - down_projection=projection == "w2", - ) - local_expert_id = global_expert_id - self.local_expert_start - logical_projection = {"w1": "gate", "w2": "down", "w3": "up"}[projection] - self.provider.load_expert_projection( - self.op_spec, - local_expert_id=local_expert_id, - projection=logical_projection, - loaded_weight=loaded_weight, - loaded_scale=loaded_scale, - w13_weight=self.w13_weight.data, - w2_weight=self.w2_weight.data, - w13_scale_inv=self.w13_scale_inv, - w2_scale_inv=self.w2_scale_inv, - ) - self._loaded_expert_shards.add(load_key) - - def validate_loaded_weights(self) -> None: - expected = { - (expert_id, projection) - for expert_id in range(self.local_expert_start, self.local_expert_end) - for projection in ("w1", "w2", "w3") - } - missing = sorted(expected - self._loaded_expert_shards) - if missing: - raise ValueError( - "Missing local MiniMax expert weights/scales: " - f"local_range=[{self.local_expert_start}, {self.local_expert_end}), " - f"missing={missing[:8]}." - ) - - def forward( - self, - hidden_states: torch.Tensor, - topk_ids: torch.Tensor, - topk_weights: torch.Tensor, - ) -> torch.Tensor: - return self.provider.run( - self.op_spec, - hidden_states, - topk_ids, - topk_weights, - self.w13_weight, - self.w2_weight, - self.w13_scale_inv, - self.w2_scale_inv, - local_expert_start=self.local_expert_start, - ep_rank=self.ep_rank, + model_label="MiniMax", + provider_resolver=resolve_moe_provider, + parallel_context=get_parallel_context(), ) diff --git a/tests/test_minimax_m2_model.py b/tests/test_minimax_m2_model.py index ba3da563..058daf75 100644 --- a/tests/test_minimax_m2_model.py +++ b/tests/test_minimax_m2_model.py @@ -487,7 +487,7 @@ def test_local_expert_loader_rejects_missing_duplicate_and_bad_tensors(): experts.load_expert_weight(0, "w1", weight, None) with pytest.raises(TypeError, match="FP8 E4M3"): experts.load_expert_weight(0, "w1", weight.float(), scale) - with pytest.raises(TypeError, match="must be FP32"): + with pytest.raises(TypeError, match=r"must use torch\.float32"): experts.load_expert_weight(0, "w1", weight, scale.bfloat16()) with pytest.raises(ValueError, match="shape mismatch"): experts.load_expert_weight(0, "w1", _random_fp8((128, 256)), scale) diff --git a/tests/test_packed_moe.py b/tests/test_packed_moe.py index 1d3247df..e5d3bb11 100644 --- a/tests/test_packed_moe.py +++ b/tests/test_packed_moe.py @@ -3,6 +3,7 @@ import torch from sparsevllm.layers.packed_moe import PackedMoeExperts +from sparsevllm.models.minimax_m2 import MiniMaxM2PackedExperts from sparsevllm.models.qwen3_moe import Qwen3MoePackedExperts from sparsevllm.operators.moe import TritonMoeProvider @@ -34,8 +35,9 @@ def _experts(**overrides) -> PackedMoeExperts: return PackedMoeExperts(**values) -def test_qwen_packed_experts_use_shared_physical_module() -> None: +def test_model_packed_experts_use_shared_physical_module() -> None: assert issubclass(Qwen3MoePackedExperts, PackedMoeExperts) + assert issubclass(MiniMaxM2PackedExperts, PackedMoeExperts) def test_packed_experts_accept_glm_router_contract_without_owning_router() -> None: From 3906ec54bd65bd10ce1bbbb8e218c67618568691 Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Wed, 12 Aug 2026 22:06:37 +0800 Subject: [PATCH 37/48] refactor: simplify glm runtime --- src/sparsevllm/layers/mla_attention.py | 138 +++++++ src/sparsevllm/models/glm4_moe_lite.py | 505 +++++-------------------- src/sparsevllm/models/minimax_m2.py | 37 +- src/sparsevllm/operators/all_reduce.py | 35 ++ src/sparsevllm/operators/moe.py | 45 ++- tests/test_glm4_moe_lite.py | 72 +--- tests/test_glm_cuda_graph.py | 1 - tests/test_mla_attention_layer.py | 18 +- tests/test_operator_providers.py | 26 ++ 9 files changed, 351 insertions(+), 526 deletions(-) diff --git a/src/sparsevllm/layers/mla_attention.py b/src/sparsevllm/layers/mla_attention.py index 2a7eac5c..d99fe078 100644 --- a/src/sparsevllm/layers/mla_attention.py +++ b/src/sparsevllm/layers/mla_attention.py @@ -10,6 +10,7 @@ AttentionViewMeta, DecodeComputeView, ExplicitKVPayload, + MlaLatentWrite, MlaLatentPayload, PrefillComputeView, ) @@ -303,6 +304,7 @@ def __init__( self.prefill_backend = TritonAttentionBackend() self._prefill_plan: _MlaPrefillPlan | None = None self._prefill_query_plan: _MlaPrefillQueryPlan | None = None + self._key_materializer_bindings: dict[int, tuple[object, Callable]] = {} @classmethod def bind( @@ -876,6 +878,142 @@ def run_decode( valid_batch_size=valid_batch_size, ) + def _ensure_key_materializer( + self, + cache_manager, + layer_idx: int, + project_latent: Callable[[torch.Tensor], torch.Tensor], + ) -> None: + layer_idx = int(layer_idx) + binding = self._key_materializer_bindings.get(layer_idx) + if binding is not None and binding[0] is cache_manager: + return + + def materialize(view: AttentionKeyComputeView) -> torch.Tensor: + return self.materialize_expanded_keys( + view, + project_latent=project_latent, + ) + + cache_manager.register_attention_key_materializer(layer_idx, materialize) + self._key_materializer_bindings[layer_idx] = (cache_manager, materialize) + + def run_cached_attention( + self, + q: torch.Tensor, + q_nope: torch.Tensor, + q_rope: torch.Tensor, + latent: torch.Tensor, + rope: torch.Tensor, + *, + project_latent: Callable[[torch.Tensor], torch.Tensor], + absorb_query: Callable[[torch.Tensor], torch.Tensor], + reconstruct_values: Callable[[torch.Tensor], torch.Tensor], + ) -> torch.Tensor: + """Execute MLA over the active cache view for the current layer.""" + + context = get_context() + cache_manager = context.cache_manager + sparse_controller = context.sparse_controller + layer_idx = int(context.now_layer_idx) + self._ensure_key_materializer(cache_manager, layer_idx, project_latent) + slot_mapping = cache_manager.store_attention_payload( + layer_idx, + MlaLatentWrite(latent=latent.unsqueeze(1), rope=rope.unsqueeze(1)), + ) + cache_manager.on_kv_stored(layer_idx, latent, slot_mapping) + + temp_slots = None + try: + if context.is_prefill: + selection = sparse_controller.get_prefill_selection(layer_idx) + cache_manager.before_prefill_layer_attention(layer_idx, selection) + view = cache_manager.build_prefill_compute_view( + layer_idx, + latent, + rope, + selection, + ) + temp_slots = view.meta.temp_slots + if context.cu_seqlens_q is None or context.cu_seqlens_q.numel() <= 1: + return torch.empty_like(q) + history = self.prepare_prefill_history( + view, + query_tokens=int(q.shape[0]), + ) + qk_nope_head_dim = self.spec.qk_head_dim - self.spec.rope_dim + projected_width = qk_nope_head_dim + self.spec.value_head_dim + expanded = project_latent(history.gathered_latent).view( + history.visible_tokens, + self.spec.local_q_heads, + projected_width, + ) + expanded_k_nope, expanded_v = expanded.split( + [qk_nope_head_dim, self.spec.value_head_dim], + dim=-1, + ) + expanded_k = torch.empty( + ( + history.visible_tokens, + self.spec.local_q_heads, + self.spec.qk_head_dim, + ), + dtype=expanded.dtype, + device=expanded.device, + ) + expanded_k[..., :qk_nope_head_dim].copy_(expanded_k_nope) + expanded_k[..., qk_nope_head_dim:].copy_( + history.gathered_rope[:, None, :] + ) + workset = self.bind_prefill_kv( + history, + expanded_k=expanded_k, + expanded_v=expanded_v, + ) + b_start_loc = context.cu_seqlens_q[:-1] + chunk_lens = context.cu_seqlens_q[1:] - context.cu_seqlens_q[:-1] + output = self.run_prefill( + q, + workset, + b_start_loc=b_start_loc, + chunk_lens=chunk_lens, + ) + explicit_view = self.build_prefill_explicit_view(workset) + cache_manager.collect_prefill_attention_score( + layer_idx, + q, + explicit_view, + b_start_loc=b_start_loc, + chunk_lens=chunk_lens, + ) + cache_manager.record_prefill_query( + layer_idx, + q, + view, + b_start_loc=b_start_loc, + chunk_lens=chunk_lens, + ) + else: + selection = sparse_controller.get_decode_selection(layer_idx, q) + view = cache_manager.build_decode_compute_view( + layer_idx, + q, + selection, + num_heads=self.spec.local_q_heads, + num_kv_heads=1, + ) + output = reconstruct_values( + self.run_decode(absorb_query(q_nope), q_rope, view) + ) + cache_manager.record_decode_query(layer_idx, q) + + sparse_controller.on_layer_attention_end(layer_idx) + cache_manager.on_layer_attention_end(layer_idx) + return output + finally: + if temp_slots is not None and temp_slots.numel() > 0: + cache_manager.release_layer_temp_slots(layer_idx, temp_slots) + __all__ = [ "MLAAttention", "MlaPrefillHistory", diff --git a/src/sparsevllm/models/glm4_moe_lite.py b/src/sparsevllm/models/glm4_moe_lite.py index 123379c1..e041fd77 100644 --- a/src/sparsevllm/models/glm4_moe_lite.py +++ b/src/sparsevllm/models/glm4_moe_lite.py @@ -9,12 +9,11 @@ from torch import nn from transformers import Glm4MoeLiteConfig -from sparsevllm.distributed import ParallelContext, ParallelGroup, get_parallel_context -from sparsevllm.models.layout import resolve_attention_qk_head_dim -from sparsevllm.engine.cache_manager import ( - AttentionKeyComputeView, - MlaLatentWrite, +from sparsevllm.distributed import ( + ParallelContext, + get_parallel_context, ) +from sparsevllm.models.layout import resolve_attention_qk_head_dim from sparsevllm.layers.embed_head import ( ParallelLMHead, VocabParallelEmbedding, @@ -32,15 +31,16 @@ from sparsevllm.models.qwen3 import Qwen3MLP from sparsevllm.operators.mla_attention import MlaAttentionOpSpec from sparsevllm.operators.all_reduce import ( - AllReduceOpSpec, PreparedAllReduceOp, - prepare_all_reduce_op, + prepare_parallel_all_reduce, ) from sparsevllm.operators.activation import resolve_silu_and_mul_provider from sparsevllm.operators.moe import ( MoeOpSpec, + append_shared_expert_route, model_activation_dtype, resolve_moe_provider, + use_packed_shared_experts, ) from sparsevllm.operators.moe_router import ( MoeRouterOpSpec, @@ -48,7 +48,6 @@ ) from sparsevllm.platforms import device_runtime from sparsevllm.utils.context import get_context -from sparsevllm.utils.log import logger from sparsevllm.utils.weight_target import WeightTarget @@ -64,8 +63,6 @@ r"^model\.layers\.(\d+)\.mlp\.shared_experts\." r"(gate_proj|up_proj|down_proj)\.weight$" ) -_MTP_LAYER_INDEX = 47 -_MTP_PREFIX = f"model.layers.{_MTP_LAYER_INDEX}." @dataclass @@ -83,35 +80,6 @@ def close(self) -> None: self._closed = True -def _prepare_decode_all_reduce( - group: ParallelGroup, - *, - max_rows: int, - hidden_size: int, - dtype: torch.dtype, - cuda_graph: bool, - device_index: int, -) -> PreparedAllReduceOp: - return prepare_all_reduce_op( - AllReduceOpSpec( - world_size=group.size, - ranks=group.ranks, - max_rows=max_rows, - hidden_size=hidden_size, - dtype=dtype, - cuda_graph=cuda_graph, - backend=( - "none" - if group.process_group is None - else str(torch.distributed.get_backend(group.process_group)) - ), - ), - group=group.process_group, - rank=group.rank, - device_index=device_index, - ) - - def build_glm4_moe_lite_runtime_config( config: Glm4MoeLiteConfig, parallel_context: ParallelContext, @@ -121,7 +89,7 @@ def build_glm4_moe_lite_runtime_config( device_index: int, ) -> Glm4MoeLiteRuntimeConfig: max_decode_tokens = int(max_decode_tokens) - moe_op = _prepare_decode_all_reduce( + moe_op = prepare_parallel_all_reduce( parallel_context.world, max_rows=2 * max_decode_tokens, hidden_size=int(config.hidden_size), @@ -132,7 +100,7 @@ def build_glm4_moe_lite_runtime_config( attention_op = ( moe_op if parallel_context.attention.ranks == parallel_context.world.ranks - else _prepare_decode_all_reduce( + else prepare_parallel_all_reduce( parallel_context.attention, max_rows=max_decode_tokens, hidden_size=int(config.hidden_size), @@ -250,7 +218,6 @@ def __init__( bias=bool(config.attention_bias), reduce_results=runtime_config is None, ) - self._attention_key_materializer_binding: tuple[int, int] | None = None def _project_kv_history(self, latent: torch.Tensor) -> torch.Tensor: if int(latent.shape[0]) <= self.proj_chunk_size: @@ -300,138 +267,6 @@ def _reconstruct_decode_values(self, latent_output: torch.Tensor) -> torch.Tenso v_weight.transpose(1, 2), ).transpose(0, 1) - def _materialize_attention_keys( - self, - view: AttentionKeyComputeView, - ) -> torch.Tensor: - return self.mla_attention.materialize_expanded_keys( - view, - project_latent=self._project_kv_history, - ) - - def _ensure_attention_key_materializer(self, cache_manager, layer_idx: int) -> None: - binding = (id(cache_manager), int(layer_idx)) - if self._attention_key_materializer_binding == binding: - return - cache_manager.register_attention_key_materializer( - int(layer_idx), - self._materialize_attention_keys, - ) - self._attention_key_materializer_binding = binding - - def _run_attention( - self, - q: torch.Tensor, - q_nope: torch.Tensor, - q_rope: torch.Tensor, - latent: torch.Tensor, - rope: torch.Tensor, - ) -> torch.Tensor: - context = get_context() - cache_manager = context.cache_manager - sparse_controller = context.sparse_controller - layer_idx = int(context.now_layer_idx) - self._ensure_attention_key_materializer(cache_manager, layer_idx) - slot_mapping = cache_manager.store_attention_payload( - layer_idx, - MlaLatentWrite( - latent=latent.unsqueeze(1), - rope=rope.unsqueeze(1), - ), - ) - cache_manager.on_kv_stored(layer_idx, latent, slot_mapping) - - temp_slots = None - try: - if context.is_prefill: - selection = sparse_controller.get_prefill_selection(layer_idx) - cache_manager.before_prefill_layer_attention(layer_idx, selection) - view = cache_manager.build_prefill_compute_view( - layer_idx, - latent, - rope, - selection, - ) - temp_slots = view.meta.temp_slots - if context.cu_seqlens_q is None or context.cu_seqlens_q.numel() <= 1: - return torch.empty_like(q) - history = self.mla_attention.prepare_prefill_history( - view, - query_tokens=int(q.shape[0]), - ) - expanded = self._project_kv_history(history.gathered_latent).view( - history.visible_tokens, - self.local_heads, - self.qk_nope_head_dim + self.v_head_dim, - ) - expanded_k_nope, expanded_v = expanded.split( - [self.qk_nope_head_dim, self.v_head_dim], - dim=-1, - ) - expanded_k = torch.empty( - ( - history.visible_tokens, - self.local_heads, - self.qk_head_dim, - ), - dtype=expanded_k_nope.dtype, - device=expanded_k_nope.device, - ) - expanded_k[..., : self.qk_nope_head_dim].copy_(expanded_k_nope) - expanded_k[..., self.qk_nope_head_dim :].copy_( - history.gathered_rope[:, None, :] - ) - workset = self.mla_attention.bind_prefill_kv( - history, - expanded_k=expanded_k, - expanded_v=expanded_v, - ) - b_start_loc = context.cu_seqlens_q[:-1] - chunk_lens = context.cu_seqlens_q[1:] - context.cu_seqlens_q[:-1] - output = self.mla_attention.run_prefill( - q, - workset, - b_start_loc=b_start_loc, - chunk_lens=chunk_lens, - ) - cache_manager.collect_prefill_attention_score( - layer_idx, - q, - self.mla_attention.build_prefill_explicit_view(workset), - b_start_loc=b_start_loc, - chunk_lens=chunk_lens, - ) - cache_manager.record_prefill_query( - layer_idx, - q, - view, - b_start_loc=b_start_loc, - chunk_lens=chunk_lens, - ) - else: - selection = sparse_controller.get_decode_selection(layer_idx, q) - view = cache_manager.build_decode_compute_view( - layer_idx, - q, - selection, - num_heads=self.local_heads, - num_kv_heads=1, - ) - latent_output = self.mla_attention.run_decode( - self._decode_absorbed_query(q_nope), - q_rope, - view, - ) - output = self._reconstruct_decode_values(latent_output) - cache_manager.record_decode_query(layer_idx, q) - - sparse_controller.on_layer_attention_end(layer_idx) - cache_manager.on_layer_attention_end(layer_idx) - return output - finally: - if temp_slots is not None and temp_slots.numel() > 0: - cache_manager.release_layer_temp_slots(layer_idx, temp_slots) - def forward( self, positions: torch.Tensor, @@ -462,12 +297,15 @@ def forward( ) k_rope = k_rope.squeeze(1) q = torch.cat((q_nope, q_rope), dim=-1) - value_output = self._run_attention( + value_output = self.mla_attention.run_cached_attention( q, q_nope, q_rope, latent, k_rope, + project_latent=self._project_kv_history, + absorb_query=self._decode_absorbed_query, + reconstruct_values=self._reconstruct_decode_values, ) output = self._project_output(value_output, hidden_states) if self.runtime_config is None: @@ -528,18 +366,15 @@ def __init__( parallel_context = get_parallel_context() self.routed_num_experts = int(config.n_routed_experts) self.routed_top_k = int(config.num_experts_per_tok) - self.fuses_shared_decode = bool( - decode_cuda_graph - and ( - self.routed_num_experts, - int(config.n_shared_experts), - self.routed_top_k, - int(config.hidden_size), - int(config.moe_intermediate_size), - int(parallel_context.moe_tp_size), - int(parallel_context.ep_size), - ) - == (64, 1, 4, 2048, 1536, 2, 1) + self.fuses_shared_decode = use_packed_shared_experts( + num_routed_experts=self.routed_num_experts, + num_shared_experts=int(config.n_shared_experts), + top_k=self.routed_top_k, + hidden_size=int(config.hidden_size), + intermediate_size=int(config.moe_intermediate_size), + tp_size=int(parallel_context.moe_tp_size), + ep_size=int(parallel_context.ep_size), + cuda_graph=decode_cuda_graph, ) packed_num_experts = self.routed_num_experts + int( self.fuses_shared_decode @@ -628,8 +463,6 @@ def forward_routed_and_shared( ) -> torch.Tensor: if self.shared_expert_id is None: raise RuntimeError("Packed shared expert is not enabled.") - from sparsevllm.kernels.triton.moe import append_shared_expert_route - fused_ids, fused_weights = append_shared_expert_route( topk_ids, topk_weights, @@ -1027,41 +860,6 @@ def forward( return hidden_states -def _expected_mtp_weight_names(num_experts: int) -> set[str]: - prefix = _MTP_PREFIX - names = { - prefix + suffix - for suffix in ( - "eh_proj.weight", - "embed_tokens.weight", - "enorm.weight", - "hnorm.weight", - "input_layernorm.weight", - "post_attention_layernorm.weight", - "self_attn.kv_a_layernorm.weight", - "self_attn.kv_a_proj_with_mqa.weight", - "self_attn.kv_b_proj.weight", - "self_attn.o_proj.weight", - "self_attn.q_a_layernorm.weight", - "self_attn.q_a_proj.weight", - "self_attn.q_b_proj.weight", - "mlp.gate.e_score_correction_bias", - "mlp.gate.weight", - "mlp.shared_experts.down_proj.weight", - "mlp.shared_experts.gate_proj.weight", - "mlp.shared_experts.up_proj.weight", - "shared_head.head.weight", - "shared_head.norm.weight", - ) - } - names.update( - f"{prefix}mlp.experts.{expert_id}.{projection}.weight" - for expert_id in range(int(num_experts)) - for projection in ("gate_proj", "up_proj", "down_proj") - ) - return names - - class Glm4MoeLiteForCausalLM(nn.Module): special_weight_loaders = (".expert_weight",) packed_modules_mapping = { @@ -1095,7 +893,6 @@ def build_runtime_kwargs( ), "mlp_chunk_size": engine_config.mlp_chunk_size, "decode_cuda_graph": decode_cuda_graph, - "expect_mtp_weights": not engine_config.tiny_random, } if parallel_context.world_size > 1: kwargs["runtime_config"] = build_glm4_moe_lite_runtime_config( @@ -1114,7 +911,6 @@ def __init__( mla_attention: MLAAttention, mlp_chunk_size: int, decode_cuda_graph: bool, - expect_mtp_weights: bool, runtime_config: Glm4MoeLiteRuntimeConfig | None = None, ) -> None: super().__init__() @@ -1131,9 +927,14 @@ def __init__( self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size) if config.tie_word_embeddings: self.lm_head.weight.data = self.model.embed_tokens.weight.data - self.expect_mtp_weights = bool(expect_mtp_weights) - self._intentionally_skipped_expert_weights: set[str] = set() - self._intentionally_skipped_mtp_weights: set[str] = set() + self.ignored_weight_prefixes = tuple( + f"model.layers.{layer_idx}." + for layer_idx in range( + int(config.num_hidden_layers), + int(config.num_hidden_layers) + + int(getattr(config, "num_nextn_predict_layers", 0)), + ) + ) def close_runtime_operators(self) -> None: if self.runtime_config is not None: @@ -1176,58 +977,11 @@ def iter_tiny_reference_weights( if source_name.endswith(down_suffix): prefix = source_name[: -len("down_proj")] for expert_id in range(int(weight.shape[0])): - yield ( - f"{prefix}{expert_id}.down_proj.weight", - weight[expert_id], - ) + yield f"{prefix}{expert_id}.down_proj.weight", weight[expert_id] continue yield source_name, weight - @torch.inference_mode() - def warmup_moe(self, num_tokens: int = 1) -> None: - num_tokens = int(num_tokens) - if num_tokens <= 0: - raise ValueError(f"num_tokens must be positive, got {num_tokens}.") - block = next( - ( - layer.mlp - for layer in self.model.layers - if isinstance(layer.mlp, Glm4MoeLiteSparseMoeBlock) - ), - None, - ) - if block is None: - raise RuntimeError("GLM model has no sparse MoE layer to warm up.") - experts = block.experts - hidden_states = torch.zeros( - (num_tokens, experts.hidden_size), - dtype=model_activation_dtype(self.config), - device=experts.w13_weight.device, - ) - block.gate(hidden_states) - top_k = int(self.config.num_experts_per_tok) - topk_ids = ( - torch.arange( - num_tokens * top_k, - dtype=torch.int64, - device=hidden_states.device, - ) - .remainder(experts.routed_num_experts) - .add(experts.local_expert_start) - .view(num_tokens, top_k) - ) - topk_weights = torch.full( - (num_tokens, top_k), - 1.0 / top_k, - dtype=torch.float32, - device=hidden_states.device, - ) - experts(hidden_states, topk_ids, topk_weights) - device_runtime.synchronize() - def map_weight_name(self, source_weight_name: str) -> str | None: - if source_weight_name.startswith(_MTP_PREFIX): - return None shared_match = _SHARED_EXPERT_SOURCE_RE.match(source_weight_name) if shared_match is not None: layer_idx, projection = shared_match.groups() @@ -1241,10 +995,8 @@ def map_weight_name(self, source_weight_name: str) -> str | None: if match is None: return source_weight_name layer_idx, global_expert_id, projection = match.groups() - block = self._sparse_block(int(layer_idx)) - global_expert_id = int(global_expert_id) - if not block.experts.is_local_expert(global_expert_id): - self._intentionally_skipped_expert_weights.add(source_weight_name) + experts = self._sparse_block(int(layer_idx)).experts + if not experts.is_local_expert(int(global_expert_id)): return None return ( f"model.layers.{layer_idx}.mlp.experts.{global_expert_id}." @@ -1259,72 +1011,10 @@ def resolve_special_weight( if match is None: return None layer_idx, global_expert_id, projection = match.groups() - experts = self._sparse_block(int(layer_idx)).experts - return WeightTarget(experts, (int(global_expert_id), projection)) - - def record_skipped_weight( - self, - source_weight_name: str, - loaded_weight_shape: tuple[int, ...] | None, - loaded_weight_dtype: str | None, - loaded_scale_shape: tuple[int, ...] | None, - loaded_scale_dtype: str | None, - ) -> None: - if source_weight_name.startswith(_MTP_PREFIX): - if loaded_weight_shape is None or any( - int(dim) <= 0 for dim in loaded_weight_shape - ): - raise ValueError( - f"Invalid GLM MTP tensor shape for {source_weight_name!r}: " - f"{loaded_weight_shape}." - ) - if loaded_weight_dtype not in {"BF16", "F32"}: - raise TypeError( - f"Unsupported GLM MTP dtype for {source_weight_name!r}: " - f"{loaded_weight_dtype}." - ) - if loaded_scale_shape is not None or loaded_scale_dtype is not None: - raise ValueError( - f"Unquantized GLM MTP tensor has an unexpected scale: " - f"{source_weight_name!r}." - ) - self._intentionally_skipped_mtp_weights.add(source_weight_name) - return - - match = _EXPERT_SOURCE_RE.match(source_weight_name) - if match is None: - raise ValueError( - f"GLM loader unexpectedly skipped {source_weight_name!r}." - ) - layer_idx, global_expert_id, projection = match.groups() - experts = self._sparse_block(int(layer_idx)).experts - global_expert_id = int(global_expert_id) - if experts.is_local_expert(global_expert_id): - raise ValueError( - f"GLM loader skipped local expert weight {source_weight_name!r}." - ) - expected_shape = ( - (experts.hidden_size, experts.global_intermediate_size) - if projection == "down_proj" - else (experts.global_intermediate_size, experts.hidden_size) + return WeightTarget( + self._sparse_block(int(layer_idx)).experts, + (int(global_expert_id), projection), ) - if loaded_weight_shape != expected_shape: - raise ValueError( - "Remote GLM expert weight shape mismatch for " - f"{source_weight_name!r}: expected={expected_shape}, " - f"got={loaded_weight_shape}." - ) - if loaded_weight_dtype not in {"BF16", "F16", "F32"}: - raise TypeError( - f"Unsupported remote GLM expert dtype for " - f"{source_weight_name!r}: {loaded_weight_dtype}." - ) - if loaded_scale_shape is not None or loaded_scale_dtype is not None: - raise ValueError( - f"Unquantized remote GLM expert has an unexpected scale: " - f"{source_weight_name!r}." - ) - self._intentionally_skipped_expert_weights.add(source_weight_name) def load_special_weight( self, @@ -1335,9 +1025,9 @@ def load_special_weight( target = self.resolve_special_weight(target_weight_name) if target is None: return 0 - global_expert_id, projection = target.shard_id + expert_id, projection = target.shard_id target.module.load_expert_weight( - global_expert_id, + expert_id, projection, loaded_weight, loaded_scale, @@ -1345,84 +1035,63 @@ def load_special_weight( return 1 def validate_loaded_weights(self, loaded_parameter_names: set[str]) -> None: - packed_expert_parameters = { + packed_experts = { name for name, _ in self.named_parameters() - if name.endswith(".mlp.experts.w13_weight") - or name.endswith(".mlp.experts.w2_weight") + if name.endswith((".mlp.experts.w13_weight", ".mlp.experts.w2_weight")) } - expected_dense_parameters = { - name for name, _ in self.named_parameters() - } - packed_expert_parameters - missing_dense = sorted(expected_dense_parameters - loaded_parameter_names) - if missing_dense: - raise ValueError( - f"Missing replicated GLM base weights: {missing_dense[:8]}." - ) - - sparse_blocks = [ - layer.mlp - for layer in self.model.layers - if isinstance(layer.mlp, Glm4MoeLiteSparseMoeBlock) - ] - for block in sparse_blocks: - block.experts.validate_loaded_weights() + missing = sorted( + {name for name, _ in self.named_parameters()} + - packed_experts + - loaded_parameter_names + ) + if missing: + raise ValueError(f"Missing replicated GLM base weights: {missing[:8]}.") + for layer_idx, layer_type in enumerate(self.config.mlp_layer_types): + if layer_type == "sparse": + self._sparse_block(layer_idx).experts.validate_loaded_weights() - expected_remote = { - f"model.layers.{layer_idx}.mlp.experts.{expert_id}." - f"{projection}.weight" - for layer_idx, layer in enumerate(self.model.layers) - if isinstance(layer.mlp, Glm4MoeLiteSparseMoeBlock) - for expert_id in range(int(self.config.n_routed_experts)) - if not layer.mlp.experts.is_local_expert(expert_id) - for projection in ("gate_proj", "up_proj", "down_proj") - } - missing_remote = sorted( - expected_remote - self._intentionally_skipped_expert_weights + @torch.inference_mode() + def warmup_moe(self, num_tokens: int = 1) -> None: + num_tokens = int(num_tokens) + if num_tokens <= 0: + raise ValueError(f"num_tokens must be positive, got {num_tokens}.") + block = next( + ( + layer.mlp + for layer in self.model.layers + if isinstance(layer.mlp, Glm4MoeLiteSparseMoeBlock) + ), + None, ) - unexpected_remote = sorted( - self._intentionally_skipped_expert_weights - expected_remote + if block is None: + raise RuntimeError("GLM model has no sparse MoE layer to warm up.") + experts = block.experts + hidden_states = torch.zeros( + (num_tokens, experts.hidden_size), + dtype=model_activation_dtype(self.config), + device=experts.w13_weight.device, ) - if missing_remote or unexpected_remote: - raise ValueError( - "GLM remote expert skip set is inconsistent: " - f"missing={missing_remote[:4]} unexpected={unexpected_remote[:4]}." - ) - - expected_mtp = _expected_mtp_weight_names(self.config.n_routed_experts) - if self.expect_mtp_weights: - missing_mtp = sorted( - expected_mtp - self._intentionally_skipped_mtp_weights - ) - unexpected_mtp = sorted( - self._intentionally_skipped_mtp_weights - expected_mtp - ) - if missing_mtp or unexpected_mtp: - raise ValueError( - "GLM MTP skip set is inconsistent: " - f"missing={missing_mtp[:4]} unexpected={unexpected_mtp[:4]}." - ) - elif self._intentionally_skipped_mtp_weights: - raise ValueError( - "GLM checkpoint supplied MTP weights when they were not expected: " - f"{sorted(self._intentionally_skipped_mtp_weights)[:4]}." + block.gate(hidden_states) + top_k = int(self.config.num_experts_per_tok) + topk_ids = ( + torch.arange( + num_tokens * top_k, + dtype=torch.int64, + device=hidden_states.device, ) - - first_sparse = sparse_blocks[0] if sparse_blocks else None - logger.info( - "Loaded GLM-4.7-Flash rank {} MLA provider={} MoE provider={} " - "attention TP {}/{} across {} base layers; local expert shards={} " - "remote expert skips={} MTP skips={}.", - self.parallel_context.world_rank, - self.model.mla_attention.provider.name, - first_sparse.experts.provider.name if first_sparse is not None else "none", - self.parallel_context.attention_tp_rank, - self.parallel_context.attention_tp_size, - len(self.model.layers), - sum(len(block.experts._loaded_expert_shards) for block in sparse_blocks), - len(self._intentionally_skipped_expert_weights), - len(self._intentionally_skipped_mtp_weights), + .remainder(experts.routed_num_experts) + .add(experts.local_expert_start) + .view(num_tokens, top_k) ) + topk_weights = torch.full( + (num_tokens, top_k), + 1.0 / top_k, + dtype=torch.float32, + device=hidden_states.device, + ) + experts(hidden_states, topk_ids, topk_weights) + device_runtime.synchronize() def forward( self, diff --git a/src/sparsevllm/models/minimax_m2.py b/src/sparsevllm/models/minimax_m2.py index 7d9a8067..1153a16f 100644 --- a/src/sparsevllm/models/minimax_m2.py +++ b/src/sparsevllm/models/minimax_m2.py @@ -10,7 +10,6 @@ from sparsevllm.distributed import ( ParallelContext, - ParallelGroup, get_parallel_context, ) from sparsevllm.layers.attention import Attention @@ -25,9 +24,8 @@ from sparsevllm.models.qwen3 import Qwen3ModelBase from sparsevllm.operators.moe import model_activation_dtype, resolve_moe_provider from sparsevllm.operators.all_reduce import ( - AllReduceOpSpec, PreparedAllReduceOp, - prepare_all_reduce_op, + prepare_parallel_all_reduce, ) from sparsevllm.operators.decode_attention import ( DecodeAttentionLaunchSpec, @@ -80,35 +78,6 @@ def close(self) -> None: self._closed = True -def _prepare_decode_all_reduce( - group: ParallelGroup, - *, - max_rows: int, - hidden_size: int, - dtype: torch.dtype, - cuda_graph: bool, - device_index: int, -) -> PreparedAllReduceOp: - return prepare_all_reduce_op( - AllReduceOpSpec( - world_size=group.size, - ranks=group.ranks, - max_rows=max_rows, - hidden_size=hidden_size, - dtype=dtype, - cuda_graph=cuda_graph, - backend=( - "none" - if group.process_group is None - else str(torch.distributed.get_backend(group.process_group)) - ), - ), - group=group.process_group, - rank=group.rank, - device_index=device_index, - ) - - def build_minimax_m2_runtime_config( config, parallel_context: ParallelContext, @@ -154,7 +123,7 @@ def build_minimax_m2_runtime_config( ), device_index=device_index, ) - moe_decode_all_reduce = _prepare_decode_all_reduce( + moe_decode_all_reduce = prepare_parallel_all_reduce( parallel_context.world, max_rows=int(max_decode_tokens), hidden_size=int(config.hidden_size), @@ -165,7 +134,7 @@ def build_minimax_m2_runtime_config( if parallel_context.attention.ranks == parallel_context.world.ranks: attention_decode_all_reduce = moe_decode_all_reduce else: - attention_decode_all_reduce = _prepare_decode_all_reduce( + attention_decode_all_reduce = prepare_parallel_all_reduce( parallel_context.attention, max_rows=int(max_decode_tokens), hidden_size=int(config.hidden_size), diff --git a/src/sparsevllm/operators/all_reduce.py b/src/sparsevllm/operators/all_reduce.py index c2314216..7acfa046 100644 --- a/src/sparsevllm/operators/all_reduce.py +++ b/src/sparsevllm/operators/all_reduce.py @@ -5,6 +5,7 @@ from dataclasses import dataclass from importlib.metadata import PackageNotFoundError, version from importlib.util import find_spec +from typing import TYPE_CHECKING import torch import torch.distributed as dist @@ -18,6 +19,9 @@ ) from sparsevllm.platforms.interface import DeviceCaps, PlatformEnum +if TYPE_CHECKING: + from sparsevllm.distributed.parallel_context import ParallelGroup + @dataclass(frozen=True) class AllReduceOpSpec: @@ -535,3 +539,34 @@ def prepare_all_reduce_op( device_index=device_index, ) return PreparedAllReduceOp(spec, provider, group=group) + + +def prepare_parallel_all_reduce( + group: ParallelGroup, + *, + max_rows: int, + hidden_size: int, + dtype: torch.dtype, + cuda_graph: bool, + device_index: int, +) -> PreparedAllReduceOp: + """Bind an all-reduce operator to an initialized parallel group.""" + + return prepare_all_reduce_op( + AllReduceOpSpec( + world_size=group.size, + ranks=group.ranks, + max_rows=max_rows, + hidden_size=hidden_size, + dtype=dtype, + cuda_graph=cuda_graph, + backend=( + "none" + if group.process_group is None + else str(dist.get_backend(group.process_group)) + ), + ), + group=group.process_group, + rank=group.rank, + device_index=device_index, + ) diff --git a/src/sparsevllm/operators/moe.py b/src/sparsevllm/operators/moe.py index 2a1148c6..31ad8679 100644 --- a/src/sparsevllm/operators/moe.py +++ b/src/sparsevllm/operators/moe.py @@ -173,9 +173,52 @@ def run( ) -> torch.Tensor: raise NotImplementedError - MOE_REGISTRY: OpRegistry[MoeOpSpec, MoeProvider] = OpRegistry("routed MoE") +_PACKED_SHARED_EXPERT_PROFILES = frozenset( + {(64, 1, 4, 2048, 1536, 2, 1)} +) + + +def use_packed_shared_experts( + *, + num_routed_experts: int, + num_shared_experts: int, + top_k: int, + hidden_size: int, + intermediate_size: int, + tp_size: int, + ep_size: int, + cuda_graph: bool, +) -> bool: + """Return whether a profiled decode path packs shared experts as routes.""" + + profile = ( + int(num_routed_experts), + int(num_shared_experts), + int(top_k), + int(hidden_size), + int(intermediate_size), + int(tp_size), + int(ep_size), + ) + return bool(cuda_graph) and profile in _PACKED_SHARED_EXPERT_PROFILES + + +def append_shared_expert_route( + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + *, + shared_expert_id: int, +) -> tuple[torch.Tensor, torch.Tensor]: + from sparsevllm.kernels.triton.moe import append_shared_expert_route as run + + return run( + topk_ids, + topk_weights, + shared_expert_id=shared_expert_id, + ) + def _sgl_moe_align_block_size( topk_ids: torch.Tensor, diff --git a/tests/test_glm4_moe_lite.py b/tests/test_glm4_moe_lite.py index a1aea9b5..c2b35284 100644 --- a/tests/test_glm4_moe_lite.py +++ b/tests/test_glm4_moe_lite.py @@ -29,7 +29,6 @@ Glm4MoeLiteForCausalLM, Glm4MoeLiteRouter, Glm4MoeLiteSparseMoeBlock, - _expected_mtp_weight_names, ) from sparsevllm.models.qwen3 import Qwen3MLP from sparsevllm.operators.mla_attention import MlaAttentionOpSpec @@ -162,7 +161,6 @@ def _model(config=None, *, tp_rank: int = 0, tp_size: int = 1): mla_attention=_fake_mla(tp_size), mlp_chunk_size=config.mlp_chunk_size, decode_cuda_graph=config.decode_cuda_graph, - expect_mtp_weights=False, ) @@ -213,7 +211,6 @@ def test_glm_runtime_kwargs_bind_model_owned_operators() -> None: "mla_attention": mla, "mlp_chunk_size": 16, "decode_cuda_graph": True, - "expect_mtp_weights": True, "runtime_config": all_reduce, } build_mla.assert_called_once_with( @@ -359,29 +356,6 @@ def test_glm_decode_absorption_and_value_reconstruction_match_linear_algebra() - torch.testing.assert_close(reconstructed, expected_reconstructed) -def test_glm_attention_binds_key_materializer_once_per_manager_context() -> None: - attention = _model().model.layers[0].self_attn - first_manager = SimpleNamespace( - register_attention_key_materializer=Mock() - ) - second_manager = SimpleNamespace( - register_attention_key_materializer=Mock() - ) - - attention._ensure_attention_key_materializer(first_manager, 0) - attention._ensure_attention_key_materializer(first_manager, 0) - attention._ensure_attention_key_materializer(second_manager, 0) - - first_manager.register_attention_key_materializer.assert_called_once_with( - 0, - attention._materialize_attention_keys, - ) - second_manager.register_attention_key_materializer.assert_called_once_with( - 0, - attention._materialize_attention_keys, - ) - - def test_glm_router_uses_bias_only_for_selection_and_scales_weights() -> None: torch.manual_seed(23) with patch( @@ -421,7 +395,6 @@ def test_tiny_transformers_weights_load_through_strict_glm_mapping() -> None: mla_attention=_fake_mla(), mlp_chunk_size=config.mlp_chunk_size, decode_cuda_graph=config.decode_cuda_graph, - expect_mtp_weights=False, ) initialize_sparse_model(model, config, seed=29) reference = build_tiny_random_hf_model(config, seed=29) @@ -449,7 +422,7 @@ def test_tiny_transformers_weights_load_through_strict_glm_mapping() -> None: @pytest.mark.parametrize(("ep_rank", "ep_size"), [(1, 2), (3, 4)]) -def test_glm_ep_loader_keeps_only_local_experts_and_accounts_remote_skips( +def test_glm_ep_loader_keeps_only_local_experts( ep_rank: int, ep_size: int, ) -> None: @@ -461,7 +434,6 @@ def test_glm_ep_loader_keeps_only_local_experts_and_accounts_remote_skips( mla_attention=_fake_mla(), mlp_chunk_size=config.mlp_chunk_size, decode_cuda_graph=config.decode_cuda_graph, - expect_mtp_weights=False, ) initialize_sparse_model(model, config, seed=31) @@ -470,9 +442,6 @@ def test_glm_ep_loader_keeps_only_local_experts_and_accounts_remote_skips( assert experts.local_expert_start == ep_rank * expected_local assert experts.local_expert_end == (ep_rank + 1) * expected_local assert len(experts._loaded_expert_shards) == expected_local * 3 - assert len(model._intentionally_skipped_expert_weights) == ( - 64 - expected_local - ) * 3 @pytest.mark.parametrize("ep_size", [1, 2, 4]) @@ -744,42 +713,3 @@ def record_broadcast(*args, **kwargs): broadcast.assert_called_once() assert broadcast.call_args.kwargs["src"] == context.expert.ranks[0] assert broadcast.call_args.kwargs["group"] is context.expert.process_group - - -def test_glm_mtp_skip_set_is_exact() -> None: - config = _config() - context = _tp_context() - with _construction_context(context): - model = Glm4MoeLiteForCausalLM( - config, - mla_attention=_fake_mla(), - mlp_chunk_size=config.mlp_chunk_size, - decode_cuda_graph=config.decode_cuda_graph, - expect_mtp_weights=True, - ) - dense_names = { - name - for name, _ in model.named_parameters() - if not name.endswith(".mlp.experts.w13_weight") - and not name.endswith(".mlp.experts.w2_weight") - } - for block in ( - layer.mlp - for layer in model.model.layers - if isinstance(layer.mlp, Glm4MoeLiteSparseMoeBlock) - ): - block.experts._loaded_expert_shards = { - (expert_id, projection) - for expert_id in range(64) - for projection in ("gate_proj", "up_proj", "down_proj") - } - expected_mtp = _expected_mtp_weight_names(64) - for name in expected_mtp: - model.record_skipped_weight(name, (1,), "BF16", None, None) - - model.validate_loaded_weights(dense_names) - assert len(model._intentionally_skipped_mtp_weights) == 212 - - model._intentionally_skipped_mtp_weights.remove(next(iter(expected_mtp))) - with pytest.raises(ValueError, match="MTP skip set is inconsistent"): - model.validate_loaded_weights(dense_names) diff --git a/tests/test_glm_cuda_graph.py b/tests/test_glm_cuda_graph.py index 7fdbe739..b8d0df5e 100644 --- a/tests/test_glm_cuda_graph.py +++ b/tests/test_glm_cuda_graph.py @@ -438,7 +438,6 @@ def _make_glm_full_graph_lane( mla_attention=mla_attention, mlp_chunk_size=8, decode_cuda_graph=True, - expect_mtp_weights=False, ) finally: torch.set_default_dtype(previous_dtype) diff --git a/tests/test_mla_attention_layer.py b/tests/test_mla_attention_layer.py index 8729b161..0d67137e 100644 --- a/tests/test_mla_attention_layer.py +++ b/tests/test_mla_attention_layer.py @@ -1,6 +1,7 @@ from __future__ import annotations -from unittest.mock import patch +from types import SimpleNamespace +from unittest.mock import Mock, patch import pytest import torch @@ -122,6 +123,21 @@ def _view( ) +def test_mla_binds_key_materializer_once_per_manager_and_layer() -> None: + attention = _attention() + project_latent = Mock() + first_manager = SimpleNamespace(register_attention_key_materializer=Mock()) + second_manager = SimpleNamespace(register_attention_key_materializer=Mock()) + + attention._ensure_key_materializer(first_manager, 0, project_latent) + attention._ensure_key_materializer(first_manager, 0, project_latent) + attention._ensure_key_materializer(first_manager, 1, project_latent) + attention._ensure_key_materializer(second_manager, 0, project_latent) + + assert first_manager.register_attention_key_materializer.call_count == 2 + second_manager.register_attention_key_materializer.assert_called_once() + + def _expand_history(attention: MLAAttention, history): heads = attention.spec.local_q_heads latent = history.gathered_latent diff --git a/tests/test_operator_providers.py b/tests/test_operator_providers.py index 9a50d8d2..58a18b9e 100644 --- a/tests/test_operator_providers.py +++ b/tests/test_operator_providers.py @@ -31,6 +31,7 @@ MoeOpSpec, SglAlignedTritonGlmMoeProvider, resolve_moe_provider, + use_packed_shared_experts, ) from sparsevllm.operators.registry import OpResolver from sparsevllm.platforms import DeviceCaps, PlatformEnum @@ -628,6 +629,31 @@ def test_glm_tp2_fused_shared_decode_uses_sgl_aligned_provider(): assert resolved.provider.name == "sgl_aligned_triton_glm" +@pytest.mark.parametrize( + ("overrides", "expected"), + [ + ({}, True), + ({"cuda_graph": False}, False), + ({"tp_size": 1}, False), + ({"num_shared_experts": 2}, False), + ], +) +def test_glm_packed_shared_expert_profile(overrides, expected): + values = { + "num_routed_experts": 64, + "num_shared_experts": 1, + "top_k": 4, + "hidden_size": 2048, + "intermediate_size": 1536, + "tp_size": 2, + "ep_size": 1, + "cuda_graph": True, + } + values.update(overrides) + + assert use_packed_shared_experts(**values) is expected + + def test_glm_tp2_ep2_moe_uses_sgl_aligned_provider(): spec = _moe_spec( activation_dtype=torch.bfloat16, From dffa9201ee9e4b0e82e26ec95d7f5d690ca868e1 Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Thu, 13 Aug 2026 15:18:30 +0800 Subject: [PATCH 38/48] docs: simplify dependency installation --- README.md | 19 +++++++------ README_zh.md | 41 ++++++---------------------- docs/en/getting_started/README.md | 24 ++++++---------- docs/zh/getting_started/README.md | 21 ++++++-------- pyproject.toml | 14 ++++++++++ requirements/locks/README.md | 20 +++++++------- tests/test_dependency_constraints.py | 23 ++++++++++++++++ 7 files changed, 83 insertions(+), 79 deletions(-) diff --git a/README.md b/README.md index ba7057fb..32455ce9 100644 --- a/README.md +++ b/README.md @@ -94,9 +94,8 @@ The full documentation index is maintained in [docs/en/README.md](docs/en/README ## Quick Start -Sparse-vLLM requires Python 3.10 or newer. The validated CUDA 12.9 runtime and -test environment is frozen in -`requirements/locks/canonical-cu129-py310.txt`. +Sparse-vLLM requires Python 3.10 or newer. Default dependencies are declared in +`pyproject.toml`. ### Conda @@ -104,9 +103,9 @@ test environment is frozen in conda create -n svllm python=3.10 -y conda activate svllm -PYTHONNOUSERSITE=1 python -s -m pip install \ - -r requirements/locks/canonical-cu129-py310.txt -PYTHONNOUSERSITE=1 python -s -m pip install --no-deps -e . +python -m pip config --site set global.extra-index-url \ + "https://download.pytorch.org/whl/cu130 https://flashinfer.ai/whl/cu130" +python -m pip install -e . # Optional MAX_JOBS=8 pip install flash-attn --no-build-isolation @@ -122,15 +121,17 @@ PyTorch wheels include their CUDA runtime, while compiled extensions such as uv venv --python 3.10 source .venv/bin/activate -uv pip install -r requirements/locks/canonical-cu129-py310.txt -uv pip install --no-deps -e . +uv pip install -e . # Optional MAX_JOBS=8 uv pip install flash-attn --no-build-isolation uv pip install flashinfer-cubin --index-url https://flashinfer.ai/whl ``` -The lock selects the CUDA 12.9 builds of PyTorch and the FlashInfer JIT cache. +uv reads the CUDA indexes from `pyproject.toml`; pip requires the one-time +environment configuration shown above. The validated CUDA 12.9 +[dependency lock](requirements/locks/README.md) is optional. + `einops`, `sgl-kernel`, and the training, benchmark, and test packages are all part of the main installation; no workflow-specific extras are required. diff --git a/README_zh.md b/README_zh.md index 021af344..f8cb4003 100644 --- a/README_zh.md +++ b/README_zh.md @@ -82,8 +82,8 @@ Sparse-vLLM 支持物理淘汰、逻辑掩码、查询感知选择和混合 KV ## 快速开始 -Sparse-vLLM 需要 Python 3.10 或更高版本。请在仓库根目录中，使用 -`pyproject.toml` 固定的运行时版本安装软件包。 +Sparse-vLLM 需要 Python 3.10 或更高版本，默认依赖声明在 +`pyproject.toml` 中。 ### Conda @@ -91,14 +91,9 @@ Sparse-vLLM 需要 Python 3.10 或更高版本。请在仓库根目录中，使 conda create -n svllm python=3.10 -y conda activate svllm -pip install torch==2.11.0 torchvision==0.26.0 triton==3.6.0 \ - --index-url https://download.pytorch.org/whl/cu130 - -# FlashInfer 在单独的索引中发布 CUDA 专用 JIT 缓存。 -pip install "flashinfer-jit-cache>=0.6.15" \ - --index-url https://flashinfer.ai/whl/cu130 - -pip install -e . +python -m pip config --site set global.extra-index-url \ + "https://download.pytorch.org/whl/cu130 https://flashinfer.ai/whl/cu130" +python -m pip install -e . # 可选安装 MAX_JOBS=8 pip install flash-attn --no-build-isolation @@ -113,10 +108,6 @@ PyTorch wheel 自带 CUDA 运行时，而 `flash-attn` 等编译扩展使用当 ```bash uv venv --python 3.10 source .venv/bin/activate -uv pip install torch==2.11.0 torchvision==0.26.0 triton==3.6.0 \ - --index-url https://download.pytorch.org/whl/cu130 -uv pip install "flashinfer-jit-cache>=0.6.15" \ - --index-url https://flashinfer.ai/whl/cu130 uv pip install -e . @@ -125,25 +116,9 @@ MAX_JOBS=8 uv pip install flash-attn --no-build-isolation uv pip install flashinfer-cubin --index-url https://flashinfer.ai/whl ``` -以上显式索引分别用于安装 CUDA 13.0 版本的 PyTorch 和 FlashInfer JIT -缓存。 - -Qwen3.5/Qwen3.6 混合注意力推理还需要安装可选 Python 依赖： - -```bash -# uv -uv pip install -e ".[qwen35]" - -# Conda/pip -pip install -e ".[qwen35]" -``` - -对于不使用 Qwen3.5/Qwen3.6 的 Vanilla、OmniKV 或 QuEST 前缀缓存卸载， -可以安装更精简的 CUDA 专用可选依赖： - -```bash -pip install -e ".[prefix-offload]" -``` +uv 会从 `pyproject.toml` 读取 CUDA index；pip 只需按上面的命令在环境中 +配置一次。已验证的 CUDA 12.9 [依赖 lock](requirements/locks/README.md) +为可选复现方式。 Qwen3.5/Qwen3.6 的 prefill causal Conv1D 和 decode Conv1D/GDN packing 路径使用仓库内置的 Triton kernel；无需安装 `sglang-kernel`，也无需编译 diff --git a/docs/en/getting_started/README.md b/docs/en/getting_started/README.md index 966db352..e7fbbdd5 100644 --- a/docs/en/getting_started/README.md +++ b/docs/en/getting_started/README.md @@ -10,9 +10,9 @@ Sparse-vLLM usage example. conda create -n svllm python=3.10 -y conda activate svllm -PYTHONNOUSERSITE=1 python -s -m pip install \ - -r requirements/locks/canonical-cu129-py310.txt -PYTHONNOUSERSITE=1 python -s -m pip install --no-deps -e . +python -m pip config --site set global.extra-index-url \ + "https://download.pytorch.org/whl/cu130 https://flashinfer.ai/whl/cu130" +python -m pip install -e . # Optional MAX_JOBS=8 pip install flash-attn --no-build-isolation @@ -20,20 +20,21 @@ MAX_JOBS=8 pip install flash-attn --no-build-isolation ## Install with uv -The canonical environment uses the CUDA 12.9 build: - ```bash uv venv --python 3.10 source .venv/bin/activate -uv pip install -r requirements/locks/canonical-cu129-py310.txt -uv pip install --no-deps -e . +uv pip install -e . # Optional MAX_JOBS=8 uv pip install flash-attn --no-build-isolation ``` +uv reads the CUDA indexes from `pyproject.toml`; pip requires the one-time +environment configuration shown above. The validated CUDA 12.9 +[dependency lock](../../../requirements/locks/README.md) is optional. + `einops`, `sgl-kernel`, and the training, benchmark, and test packages are all runtime dependencies, so workflow-specific extras are not required. @@ -43,16 +44,9 @@ block-scaled FP8 formats. Its prefill causal Conv1D and decode Conv1D/GDN packing paths use local Triton kernels and do not call `sgl-kernel` themselves. -The required `flashinfer-jit-cache` package provides modules built for a -specific CUDA toolkit version. Select the index matching the CUDA version used -by PyTorch. `flashinfer-cubin` is an optional acceleration package containing -architecture-specific device binaries: +`flashinfer-cubin` is an optional acceleration package: ```bash -pip install "flashinfer-jit-cache==0.6.15.post1" \ - --index-url https://flashinfer.ai/whl/cu129 - -# Optional pip install flashinfer-cubin --index-url https://flashinfer.ai/whl ``` diff --git a/docs/zh/getting_started/README.md b/docs/zh/getting_started/README.md index 15bbaaa0..4ef5f068 100644 --- a/docs/zh/getting_started/README.md +++ b/docs/zh/getting_started/README.md @@ -8,9 +8,9 @@ conda create -n svllm python=3.10 -y conda activate svllm -PYTHONNOUSERSITE=1 python -s -m pip install \ - -r requirements/locks/canonical-cu129-py310.txt -PYTHONNOUSERSITE=1 python -s -m pip install --no-deps -e . +python -m pip config --site set global.extra-index-url \ + "https://download.pytorch.org/whl/cu130 https://flashinfer.ai/whl/cu130" +python -m pip install -e . # Optional MAX_JOBS=8 pip install flash-attn --no-build-isolation @@ -18,19 +18,20 @@ MAX_JOBS=8 pip install flash-attn --no-build-isolation ## 使用 uv 安装 -Canonical 环境使用 CUDA 12.9 build： - ```bash uv venv --python 3.10 source .venv/bin/activate -uv pip install -r requirements/locks/canonical-cu129-py310.txt -uv pip install --no-deps -e . +uv pip install -e . # Optional MAX_JOBS=8 uv pip install flash-attn --no-build-isolation ``` +uv 会从 `pyproject.toml` 读取 CUDA index；pip 只需按上面的命令在环境中 +配置一次。已验证的 CUDA 12.9 +[依赖 lock](../../../requirements/locks/README.md) 为可选复现方式。 + `einops`、`sgl-kernel` 以及训练、benchmark 和测试包均已是主依赖， 不再需要工作流专用 extra。 @@ -39,13 +40,9 @@ Sparse-vLLM 当前支持未量化 BF16 和 block-scaled FP8 格式的 Qwen3.5/Qw 其 prefill causal Conv1D 和 decode Conv1D/GDN packing path 使用仓库本地 Triton kernel，本身不调用 `sgl-kernel`。 -必需的 `flashinfer-jit-cache` package 提供针对特定 CUDA toolkit 版本构建的 module。请选择与 PyTorch 所用 CUDA 版本匹配的 index。`flashinfer-cubin` 是可选加速 package，包含特定架构的 device binary： +`flashinfer-cubin` 是可选加速 package： ```bash -pip install "flashinfer-jit-cache==0.6.15.post1" \ - --index-url https://flashinfer.ai/whl/cu129 - -# Optional pip install flashinfer-cubin --index-url https://flashinfer.ai/whl ``` diff --git a/pyproject.toml b/pyproject.toml index b185ac09..daaeee1f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,3 +56,17 @@ include = ["sparsevllm*"] [tool.pytest.ini_options] testpaths = ["tests"] + +[tool.uv.sources] +torch = { index = "pytorch-cu130" } +flashinfer-jit-cache = { index = "flashinfer-cu130" } + +[[tool.uv.index]] +name = "pytorch-cu130" +url = "https://download.pytorch.org/whl/cu130" +explicit = true + +[[tool.uv.index]] +name = "flashinfer-cu130" +url = "https://flashinfer.ai/whl/cu130" +explicit = true diff --git a/requirements/locks/README.md b/requirements/locks/README.md index 540b7136..768f842d 100644 --- a/requirements/locks/README.md +++ b/requirements/locks/README.md @@ -1,11 +1,10 @@ # Dependency locks -`canonical-cu129-py310.txt` is the supported runtime and test baseline for -Python 3.10 and CUDA 12.9. It freezes the complete resolved Python environment; -`pyproject.toml` describes intentionally broader direct-dependency compatibility -ranges. Training, benchmark, and test packages are part of the main install. -The lock remains the only fully validated package combination; versions inside -the wider ranges are candidates until they pass the same validation gates. +`canonical-cu129-py310.txt` is an optional, fully validated reproducibility +baseline for Python 3.10 and CUDA 12.9. It freezes the complete resolved Python +environment; `pyproject.toml` is the default installation contract and describes +intentionally broader direct-dependency compatibility ranges. Training, +benchmark, and test packages are part of the main install. Create an isolated environment from the repository root: @@ -15,7 +14,8 @@ uv pip install -r requirements/locks/canonical-cu129-py310.txt uv pip install --no-deps -e . ``` -The lock is the compatibility contract. An unlocked install is a development -convenience, not evidence that a newly resolved dependency combination is -supported. Replace the lock only after dependency checks, CPU tests, focused -GPU operator tests, and a real model-path smoke all pass. +For reproducible runs, the lock is the compatibility contract. The default +unlocked install is the development path, not evidence that a newly resolved +dependency combination is fully validated. Replace the lock only after +dependency checks, CPU tests, focused GPU operator tests, and a real model-path +smoke all pass. diff --git a/tests/test_dependency_constraints.py b/tests/test_dependency_constraints.py index c7837fe4..2056d360 100644 --- a/tests/test_dependency_constraints.py +++ b/tests/test_dependency_constraints.py @@ -46,6 +46,29 @@ def test_canonical_lock_pins_validated_tilelang_runtime(): assert "apache-tvm-ffi==0.1.10" in locked_requirements +def test_uv_routes_cuda_packages_to_explicit_indexes(): + pyproject_path = Path(__file__).parents[1] / "pyproject.toml" + config = tomllib.loads(pyproject_path.read_text()) + uv_config = config["tool"]["uv"] + + assert uv_config["sources"] == { + "torch": {"index": "pytorch-cu130"}, + "flashinfer-jit-cache": {"index": "flashinfer-cu130"}, + } + assert uv_config["index"] == [ + { + "name": "pytorch-cu130", + "url": "https://download.pytorch.org/whl/cu130", + "explicit": True, + }, + { + "name": "flashinfer-cu130", + "url": "https://flashinfer.ai/whl/cu130", + "explicit": True, + }, + ] + + def test_workflow_dependencies_are_part_of_main_install(): pyproject_path = Path(__file__).parents[1] / "pyproject.toml" project = tomllib.loads(pyproject_path.read_text())["project"] From a47e481694aad7ccf7e738c1aad573d3b3ac836a Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Thu, 13 Aug 2026 15:33:19 +0800 Subject: [PATCH 39/48] feat: add selectable CUDA extras --- README.md | 10 ++++---- README_zh.md | 12 +++++----- docs/en/getting_started/README.md | 10 ++++---- docs/zh/getting_started/README.md | 10 ++++---- pyproject.toml | 35 +++++++++++++++++++++++---- tests/test_dependency_constraints.py | 36 +++++++++++++++++++++++----- 6 files changed, 81 insertions(+), 32 deletions(-) diff --git a/README.md b/README.md index 32455ce9..495a95e5 100644 --- a/README.md +++ b/README.md @@ -103,9 +103,10 @@ Sparse-vLLM requires Python 3.10 or newer. Default dependencies are declared in conda create -n svllm python=3.10 -y conda activate svllm +CUDA_VERSION=cu130 python -m pip config --site set global.extra-index-url \ - "https://download.pytorch.org/whl/cu130 https://flashinfer.ai/whl/cu130" -python -m pip install -e . + "https://download.pytorch.org/whl/${CUDA_VERSION} https://flashinfer.ai/whl/${CUDA_VERSION}" +python -m pip install -e ".[${CUDA_VERSION}]" # Optional MAX_JOBS=8 pip install flash-attn --no-build-isolation @@ -121,15 +122,14 @@ PyTorch wheels include their CUDA runtime, while compiled extensions such as uv venv --python 3.10 source .venv/bin/activate -uv pip install -e . +uv pip install -e ".[cu130]" # Optional MAX_JOBS=8 uv pip install flash-attn --no-build-isolation uv pip install flashinfer-cubin --index-url https://flashinfer.ai/whl ``` -uv reads the CUDA indexes from `pyproject.toml`; pip requires the one-time -environment configuration shown above. The validated CUDA 12.9 +Use `cu129` instead of `cu130` for CUDA 12.9. The validated CUDA 12.9 [dependency lock](requirements/locks/README.md) is optional. `einops`, `sgl-kernel`, and the training, benchmark, and test packages are all diff --git a/README_zh.md b/README_zh.md index f8cb4003..1169af80 100644 --- a/README_zh.md +++ b/README_zh.md @@ -91,9 +91,10 @@ Sparse-vLLM 需要 Python 3.10 或更高版本，默认依赖声明在 conda create -n svllm python=3.10 -y conda activate svllm +CUDA_VERSION=cu130 python -m pip config --site set global.extra-index-url \ - "https://download.pytorch.org/whl/cu130 https://flashinfer.ai/whl/cu130" -python -m pip install -e . + "https://download.pytorch.org/whl/${CUDA_VERSION} https://flashinfer.ai/whl/${CUDA_VERSION}" +python -m pip install -e ".[${CUDA_VERSION}]" # 可选安装 MAX_JOBS=8 pip install flash-attn --no-build-isolation @@ -109,16 +110,15 @@ PyTorch wheel 自带 CUDA 运行时，而 `flash-attn` 等编译扩展使用当 uv venv --python 3.10 source .venv/bin/activate -uv pip install -e . +uv pip install -e ".[cu130]" # 可选安装 MAX_JOBS=8 uv pip install flash-attn --no-build-isolation uv pip install flashinfer-cubin --index-url https://flashinfer.ai/whl ``` -uv 会从 `pyproject.toml` 读取 CUDA index；pip 只需按上面的命令在环境中 -配置一次。已验证的 CUDA 12.9 [依赖 lock](requirements/locks/README.md) -为可选复现方式。 +CUDA 12.9 环境将 `cu130` 换成 `cu129`。已验证的 CUDA 12.9 +[依赖 lock](requirements/locks/README.md) 为可选复现方式。 Qwen3.5/Qwen3.6 的 prefill causal Conv1D 和 decode Conv1D/GDN packing 路径使用仓库内置的 Triton kernel；无需安装 `sglang-kernel`，也无需编译 diff --git a/docs/en/getting_started/README.md b/docs/en/getting_started/README.md index e7fbbdd5..688d633f 100644 --- a/docs/en/getting_started/README.md +++ b/docs/en/getting_started/README.md @@ -10,9 +10,10 @@ Sparse-vLLM usage example. conda create -n svllm python=3.10 -y conda activate svllm +CUDA_VERSION=cu130 python -m pip config --site set global.extra-index-url \ - "https://download.pytorch.org/whl/cu130 https://flashinfer.ai/whl/cu130" -python -m pip install -e . + "https://download.pytorch.org/whl/${CUDA_VERSION} https://flashinfer.ai/whl/${CUDA_VERSION}" +python -m pip install -e ".[${CUDA_VERSION}]" # Optional MAX_JOBS=8 pip install flash-attn --no-build-isolation @@ -24,15 +25,14 @@ MAX_JOBS=8 pip install flash-attn --no-build-isolation uv venv --python 3.10 source .venv/bin/activate -uv pip install -e . +uv pip install -e ".[cu130]" # Optional MAX_JOBS=8 uv pip install flash-attn --no-build-isolation ``` -uv reads the CUDA indexes from `pyproject.toml`; pip requires the one-time -environment configuration shown above. The validated CUDA 12.9 +Use `cu129` instead of `cu130` for CUDA 12.9. The validated CUDA 12.9 [dependency lock](../../../requirements/locks/README.md) is optional. `einops`, `sgl-kernel`, and the training, benchmark, and test packages are all diff --git a/docs/zh/getting_started/README.md b/docs/zh/getting_started/README.md index 4ef5f068..42189a22 100644 --- a/docs/zh/getting_started/README.md +++ b/docs/zh/getting_started/README.md @@ -8,9 +8,10 @@ conda create -n svllm python=3.10 -y conda activate svllm +CUDA_VERSION=cu130 python -m pip config --site set global.extra-index-url \ - "https://download.pytorch.org/whl/cu130 https://flashinfer.ai/whl/cu130" -python -m pip install -e . + "https://download.pytorch.org/whl/${CUDA_VERSION} https://flashinfer.ai/whl/${CUDA_VERSION}" +python -m pip install -e ".[${CUDA_VERSION}]" # Optional MAX_JOBS=8 pip install flash-attn --no-build-isolation @@ -22,14 +23,13 @@ MAX_JOBS=8 pip install flash-attn --no-build-isolation uv venv --python 3.10 source .venv/bin/activate -uv pip install -e . +uv pip install -e ".[cu130]" # Optional MAX_JOBS=8 uv pip install flash-attn --no-build-isolation ``` -uv 会从 `pyproject.toml` 读取 CUDA index；pip 只需按上面的命令在环境中 -配置一次。已验证的 CUDA 12.9 +CUDA 12.9 环境将 `cu130` 换成 `cu129`。已验证的 CUDA 12.9 [依赖 lock](../../../requirements/locks/README.md) 为可选复现方式。 `einops`、`sgl-kernel` 以及训练、benchmark 和测试包均已是主依赖， diff --git a/pyproject.toml b/pyproject.toml index daaeee1f..ee454b17 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,7 +14,6 @@ authors = [ ] dependencies = [ "fire", - "torch>=2.9,<3", "transformers>=5.13,<6", "socksio>=1,<2", "triton>=3.5,<4", @@ -28,8 +27,6 @@ dependencies = [ "fastapi>=0.100", "uvicorn[standard]", "pydantic>=2", - "flashinfer-python>=0.6.15,<0.7", - "flashinfer-jit-cache>=0.6.15,<0.7", "accelerate", "datasets", "wandb", @@ -45,6 +42,18 @@ dependencies = [ "tomli; python_version < '3.11'", ] +[project.optional-dependencies] +cu129 = [ + "torch>=2.9,<3", + "flashinfer-python[cu12]>=0.6.15,<0.7", + "flashinfer-jit-cache>=0.6.15,<0.7", +] +cu130 = [ + "torch>=2.9,<3", + "flashinfer-python[cu13]>=0.6.15,<0.7", + "flashinfer-jit-cache>=0.6.15,<0.7", +] + [project.scripts] sparsevllm-openai-server = "sparsevllm.entrypoints.openai.api_server:main" sparsevllm-openai-router = "sparsevllm.entrypoints.openai.smart_router:main" @@ -58,14 +67,30 @@ include = ["sparsevllm*"] testpaths = ["tests"] [tool.uv.sources] -torch = { index = "pytorch-cu130" } -flashinfer-jit-cache = { index = "flashinfer-cu130" } +torch = [ + { index = "pytorch-cu129", extra = "cu129" }, + { index = "pytorch-cu130", extra = "cu130" }, +] +flashinfer-jit-cache = [ + { index = "flashinfer-cu129", extra = "cu129" }, + { index = "flashinfer-cu130", extra = "cu130" }, +] + +[[tool.uv.index]] +name = "pytorch-cu129" +url = "https://download.pytorch.org/whl/cu129" +explicit = true [[tool.uv.index]] name = "pytorch-cu130" url = "https://download.pytorch.org/whl/cu130" explicit = true +[[tool.uv.index]] +name = "flashinfer-cu129" +url = "https://flashinfer.ai/whl/cu129" +explicit = true + [[tool.uv.index]] name = "flashinfer-cu130" url = "https://flashinfer.ai/whl/cu130" diff --git a/tests/test_dependency_constraints.py b/tests/test_dependency_constraints.py index 2056d360..ffefdd58 100644 --- a/tests/test_dependency_constraints.py +++ b/tests/test_dependency_constraints.py @@ -11,13 +11,10 @@ def test_runtime_compatibility_bounds_cover_canonical_lock(): project = tomllib.loads(pyproject_path.read_text())["project"] dependencies = set(project["dependencies"]) - assert "torch>=2.9,<3" in dependencies assert "triton>=3.5,<4" in dependencies assert "tilelang==0.1.9" in dependencies assert "apache-tvm-ffi==0.1.10" in dependencies assert "transformers>=5.13,<6" in dependencies - assert "flashinfer-python>=0.6.15,<0.7" in dependencies - assert "flashinfer-jit-cache>=0.6.15,<0.7" in dependencies assert "sgl-kernel>=0.3.21,<0.4" in dependencies assert {"fire", "pillow", "einops", "tqdm", "loguru"} <= dependencies assert not any( @@ -51,16 +48,44 @@ def test_uv_routes_cuda_packages_to_explicit_indexes(): config = tomllib.loads(pyproject_path.read_text()) uv_config = config["tool"]["uv"] + assert config["project"]["optional-dependencies"] == { + "cu129": [ + "torch>=2.9,<3", + "flashinfer-python[cu12]>=0.6.15,<0.7", + "flashinfer-jit-cache>=0.6.15,<0.7", + ], + "cu130": [ + "torch>=2.9,<3", + "flashinfer-python[cu13]>=0.6.15,<0.7", + "flashinfer-jit-cache>=0.6.15,<0.7", + ], + } assert uv_config["sources"] == { - "torch": {"index": "pytorch-cu130"}, - "flashinfer-jit-cache": {"index": "flashinfer-cu130"}, + "torch": [ + {"index": "pytorch-cu129", "extra": "cu129"}, + {"index": "pytorch-cu130", "extra": "cu130"}, + ], + "flashinfer-jit-cache": [ + {"index": "flashinfer-cu129", "extra": "cu129"}, + {"index": "flashinfer-cu130", "extra": "cu130"}, + ], } assert uv_config["index"] == [ + { + "name": "pytorch-cu129", + "url": "https://download.pytorch.org/whl/cu129", + "explicit": True, + }, { "name": "pytorch-cu130", "url": "https://download.pytorch.org/whl/cu130", "explicit": True, }, + { + "name": "flashinfer-cu129", + "url": "https://flashinfer.ai/whl/cu129", + "explicit": True, + }, { "name": "flashinfer-cu130", "url": "https://flashinfer.ai/whl/cu130", @@ -74,7 +99,6 @@ def test_workflow_dependencies_are_part_of_main_install(): project = tomllib.loads(pyproject_path.read_text())["project"] dependencies = set(project["dependencies"]) - assert "optional-dependencies" not in project assert { "accelerate", "datasets", From f47925f4a19310cd56db4187d3259355bdd76789 Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Thu, 13 Aug 2026 16:50:47 +0800 Subject: [PATCH 40/48] fix: constrain cutlass dsl version --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index ee454b17..ab0e5c11 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,6 +19,7 @@ dependencies = [ "triton>=3.5,<4", "tilelang==0.1.9", "apache-tvm-ffi==0.1.10", + "nvidia-cutlass-dsl>=4.6,<5", "pillow", "einops", "sgl-kernel>=0.3.21,<0.4", From dfec1d14fc1e197ee689e3a2100645b98df3f9e5 Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Thu, 13 Aug 2026 17:01:06 +0800 Subject: [PATCH 41/48] fix: make cuda extras mutually exclusive --- pyproject.toml | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index ab0e5c11..7f97289c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -67,6 +67,14 @@ include = ["sparsevllm*"] [tool.pytest.ini_options] testpaths = ["tests"] +[tool.uv] +conflicts = [ + [ + { extra = "cu129" }, + { extra = "cu130" }, + ], +] + [tool.uv.sources] torch = [ { index = "pytorch-cu129", extra = "cu129" }, From 593b3efe113b6715a8bb6d071d04f0a533a64643 Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Thu, 13 Aug 2026 17:31:05 +0800 Subject: [PATCH 42/48] fix: validate sgl kernel binaries --- src/sparsevllm/kernels/external/sgl/fa3.py | 36 ++++++++------ src/sparsevllm/kernels/external/sgl/moe.py | 17 ++++++- .../kernels/external/sgl/support.py | 16 ++++++- tests/test_sgl_fa3.py | 47 +++++++++++++++++++ tests/test_sgl_moe.py | 15 ++++++ 5 files changed, 114 insertions(+), 17 deletions(-) diff --git a/src/sparsevllm/kernels/external/sgl/fa3.py b/src/sparsevllm/kernels/external/sgl/fa3.py index 45cf1f0e..db1df9c2 100644 --- a/src/sparsevllm/kernels/external/sgl/fa3.py +++ b/src/sparsevllm/kernels/external/sgl/fa3.py @@ -1,5 +1,6 @@ from __future__ import annotations +import importlib from collections.abc import Sequence from dataclasses import dataclass @@ -98,8 +99,27 @@ def matches( ) +def _sgl_fa3_op(): + # Register FA3 without importing the unrelated optional FA4 wrapper. + importlib.import_module("sgl_kernel.flash_ops") + return torch.ops.sgl_kernel.fwd.default + + def sgl_fa3_support() -> tuple[bool, str]: - return sgl_kernel_support("FA3 raw op") + supported, reason = sgl_kernel_support("FA3 raw op") + if not supported: + return supported, reason + try: + op = _sgl_fa3_op() + argument_names = tuple(argument.name for argument in op._schema.arguments) + except Exception as error: + return False, ( + "sgl-kernel FA3 raw op failed to load: " + f"{type(error).__name__}: {error}" + ) + if argument_names != _FWD_ARGUMENTS: + return False, f"unsupported sgl-kernel FA3 fwd schema: {argument_names}" + return True, reason class SglFa3DecodeKernel: @@ -117,19 +137,7 @@ def __init__( if not supported: raise RuntimeError(reason) - # Importing flash_ops registers FA3 only. Importing - # sgl_kernel.flash_attn would also import the optional FA4 wrapper, - # whose CUTLASS Python dependency is unrelated to this provider. - from sgl_kernel import flash_ops as _flash_ops # noqa: F401 - - op = torch.ops.sgl_kernel.fwd.default - argument_names = tuple(argument.name for argument in op._schema.arguments) - if argument_names != _FWD_ARGUMENTS: - raise RuntimeError( - "Unsupported sgl-kernel FA3 fwd schema: " - f"arguments={argument_names}." - ) - self._op = op + self._op = _sgl_fa3_op() scheduler_packet = getattr( torch.ops.sgl_kernel, "get_scheduler_metadata", diff --git a/src/sparsevllm/kernels/external/sgl/moe.py b/src/sparsevllm/kernels/external/sgl/moe.py index b53296af..5f710837 100644 --- a/src/sparsevllm/kernels/external/sgl/moe.py +++ b/src/sparsevllm/kernels/external/sgl/moe.py @@ -1,5 +1,7 @@ from __future__ import annotations +import importlib + import torch import triton @@ -10,7 +12,20 @@ def sgl_moe_alignment_support() -> tuple[bool, str]: """Check the SGL expert-alignment API used by the Triton MoE provider.""" - return sgl_kernel_support("MoE alignment") + supported, reason = sgl_kernel_support("MoE alignment") + if not supported: + return supported, reason + try: + alignment = importlib.import_module("sgl_kernel").moe_align_block_size + except Exception as error: + return False, ( + "sgl-kernel MoE alignment failed to load: " + f"{type(error).__name__}: {error}" + ) + return (True, reason) if callable(alignment) else ( + False, + "sgl-kernel moe_align_block_size is not callable", + ) def sgl_moe_align_block_size( diff --git a/src/sparsevllm/kernels/external/sgl/support.py b/src/sparsevllm/kernels/external/sgl/support.py index 110761f4..ba1a18b2 100644 --- a/src/sparsevllm/kernels/external/sgl/support.py +++ b/src/sparsevllm/kernels/external/sgl/support.py @@ -1,5 +1,6 @@ from __future__ import annotations +import importlib import importlib.metadata import importlib.util import re @@ -9,9 +10,13 @@ def sgl_kernel_support(feature: str) -> tuple[bool, str]: - """Check the sgl-kernel API range declared by pyproject.toml.""" + """Check the declared version and load the architecture-specific extension.""" - if importlib.util.find_spec("sgl_kernel") is None: + try: + package_spec = importlib.util.find_spec("sgl_kernel") + except (ImportError, ValueError) as error: + return False, f"sgl-kernel package discovery failed: {error}" + if package_spec is None: return False, "sgl-kernel is not installed" try: version = importlib.metadata.version("sgl-kernel") @@ -21,4 +26,11 @@ def sgl_kernel_support(feature: str) -> tuple[bool, str]: parsed = tuple(map(int, match.groups())) if match else None if parsed is None or not _MIN_VERSION <= parsed < _MAX_VERSION: return False, f"requires sgl-kernel>=0.3.21,<0.4, got {version}" + try: + importlib.import_module("sgl_kernel") + except Exception as error: + return False, ( + f"sgl-kernel {version} {feature} failed to load: " + f"{type(error).__name__}: {error}" + ) return True, f"sgl-kernel {version} {feature} is available" diff --git a/tests/test_sgl_fa3.py b/tests/test_sgl_fa3.py index 2f211877..30f28b28 100644 --- a/tests/test_sgl_fa3.py +++ b/tests/test_sgl_fa3.py @@ -1,11 +1,13 @@ from __future__ import annotations +from types import SimpleNamespace from unittest.mock import patch import pytest import torch from sparsevllm.kernels.external.sgl.fa3 import ( + _FWD_ARGUMENTS, SglFa3DecodeKernel, sgl_fa3_support, ) @@ -39,9 +41,22 @@ def test_sgl_fa3_support_rejects_04_package() -> None: def test_sgl_fa3_support_accepts_declared_package() -> None: + op = SimpleNamespace( + _schema=SimpleNamespace( + arguments=[ + SimpleNamespace(name=name) + for name in _FWD_ARGUMENTS + ] + ) + ) with ( + patch( + "sparsevllm.kernels.external.sgl.fa3._sgl_fa3_op", + return_value=op, + ), patch("importlib.util.find_spec", return_value=object()), patch("importlib.metadata.version", return_value="0.3.21"), + patch("importlib.import_module", return_value=object()), ): supported, reason = sgl_fa3_support() @@ -49,6 +64,38 @@ def test_sgl_fa3_support_accepts_declared_package() -> None: assert "0.3.21" in reason +def test_sgl_fa3_support_rejects_binary_load_failure() -> None: + with ( + patch("importlib.util.find_spec", return_value=object()), + patch("importlib.metadata.version", return_value="0.3.21"), + patch( + "importlib.import_module", + side_effect=ImportError("undefined symbol: c10_cuda_check"), + ), + ): + supported, reason = sgl_fa3_support() + + assert not supported + assert "undefined symbol" in reason + + +def test_sgl_fa3_support_rejects_missing_op_schema() -> None: + with ( + patch( + "sparsevllm.kernels.external.sgl.fa3.sgl_kernel_support", + return_value=(True, "available"), + ), + patch( + "sparsevllm.kernels.external.sgl.fa3._sgl_fa3_op", + return_value=object(), + ), + ): + supported, reason = sgl_fa3_support() + + assert not supported + assert "failed to load" in reason + + @pytest.mark.skipif( not torch.cuda.is_available() or not sgl_fa3_support()[0], reason="CUDA and a validated sgl-kernel are required", diff --git a/tests/test_sgl_moe.py b/tests/test_sgl_moe.py index 932b891b..825522ec 100644 --- a/tests/test_sgl_moe.py +++ b/tests/test_sgl_moe.py @@ -1,5 +1,6 @@ from __future__ import annotations +from types import SimpleNamespace from unittest.mock import patch import pytest @@ -33,9 +34,11 @@ def test_sgl_moe_support_rejects_04_package() -> None: def test_sgl_moe_support_accepts_declared_package() -> None: + module = SimpleNamespace(moe_align_block_size=lambda *_args: None) with ( patch("importlib.util.find_spec", return_value=object()), patch("importlib.metadata.version", return_value="0.3.21"), + patch("importlib.import_module", return_value=module), ): supported, reason = sgl_moe_alignment_support() @@ -43,6 +46,18 @@ def test_sgl_moe_support_accepts_declared_package() -> None: assert "0.3.21" in reason +def test_sgl_moe_support_rejects_missing_alignment_api() -> None: + with ( + patch("importlib.util.find_spec", return_value=object()), + patch("importlib.metadata.version", return_value="0.3.21"), + patch("importlib.import_module", return_value=SimpleNamespace()), + ): + supported, reason = sgl_moe_alignment_support() + + assert not supported + assert "moe_align_block_size" in reason + + @pytest.mark.skipif( not torch.cuda.is_available() or not sgl_moe_alignment_support()[0], reason="CUDA and a validated sgl-kernel are required", From 3839783cbe63aa3f04e84f5d19ead6e69ea80434 Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Thu, 13 Aug 2026 20:07:04 +0800 Subject: [PATCH 43/48] feat: support sglang kernel 0.4.5 --- README.md | 4 +-- docs/en/getting_started/README.md | 4 +-- docs/en/getting_started/reproducibility.md | 4 +-- docs/zh/getting_started/README.md | 4 +-- docs/zh/getting_started/reproducibility.md | 4 +-- pyproject.toml | 6 ++-- src/sparsevllm/kernels/external/sgl/fa3.py | 18 +++++----- src/sparsevllm/kernels/external/sgl/moe.py | 8 ++--- .../kernels/external/sgl/support.py | 19 +++++----- tests/test_dependency_constraints.py | 7 ++-- tests/test_mla_attention_operator.py | 6 ++-- tests/test_sgl_fa3.py | 35 +++++++------------ tests/test_sgl_moe.py | 24 +++++++------ 13 files changed, 70 insertions(+), 73 deletions(-) diff --git a/README.md b/README.md index 495a95e5..2b87ad4b 100644 --- a/README.md +++ b/README.md @@ -132,14 +132,14 @@ uv pip install flashinfer-cubin --index-url https://flashinfer.ai/whl Use `cu129` instead of `cu130` for CUDA 12.9. The validated CUDA 12.9 [dependency lock](requirements/locks/README.md) is optional. -`einops`, `sgl-kernel`, and the training, benchmark, and test packages are all +`einops`, `sglang-kernel`, and the training, benchmark, and test packages are all part of the main installation; no workflow-specific extras are required. Sparse-vLLM supports Qwen3.5/Qwen3.6 checkpoints in unquantized BF16 and block-scaled FP8 formats. The Qwen3.5/Qwen3.6 prefill causal Conv1D and decode Conv1D/GDN packing paths -use repository-local Triton kernels; they do not call `sgl-kernel` themselves. +use repository-local Triton kernels; they do not call `sglang-kernel` themselves. For the full dependency list and a minimal `LLM(...)` example, see [Getting Started](docs/en/getting_started/README.md). diff --git a/docs/en/getting_started/README.md b/docs/en/getting_started/README.md index 688d633f..b0b12550 100644 --- a/docs/en/getting_started/README.md +++ b/docs/en/getting_started/README.md @@ -35,14 +35,14 @@ MAX_JOBS=8 uv pip install flash-attn --no-build-isolation Use `cu129` instead of `cu130` for CUDA 12.9. The validated CUDA 12.9 [dependency lock](../../../requirements/locks/README.md) is optional. -`einops`, `sgl-kernel`, and the training, benchmark, and test packages are all +`einops`, `sglang-kernel`, and the training, benchmark, and test packages are all runtime dependencies, so workflow-specific extras are not required. Sparse-vLLM supports Qwen3.5/Qwen3.6 checkpoints in unquantized BF16 and block-scaled FP8 formats. Its prefill causal Conv1D and decode Conv1D/GDN packing paths use local Triton -kernels and do not call `sgl-kernel` themselves. +kernels and do not call `sglang-kernel` themselves. `flashinfer-cubin` is an optional acceleration package: diff --git a/docs/en/getting_started/reproducibility.md b/docs/en/getting_started/reproducibility.md index 1f2204c5..13504757 100644 --- a/docs/en/getting_started/reproducibility.md +++ b/docs/en/getting_started/reproducibility.md @@ -11,10 +11,10 @@ The README contains the current install command. The expected baseline is: - Python 3.10. - The complete runtime and test environment from `requirements/locks/canonical-cu129-py310.txt`. -- PyTorch 2.9.1 with the CUDA 12.9 wheel and Triton 3.5.1. +- PyTorch 2.11.0 with the matching CUDA wheel and Triton 3.6.0. - `flashinfer-python==0.6.15.post1` and the CUDA 12.9 build of `flashinfer-jit-cache==0.6.15.post1`. -- `sgl-kernel==0.3.21` and `einops>=0.8.2` as runtime dependencies. +- `sglang-kernel==0.4.5` and `einops>=0.8.2` as runtime dependencies. - Optional matching `flashinfer-cubin` installed from the generic FlashInfer wheel index when precompiled device binaries are desired. - `transformers==5.13.1`. diff --git a/docs/zh/getting_started/README.md b/docs/zh/getting_started/README.md index 42189a22..bc4f76c9 100644 --- a/docs/zh/getting_started/README.md +++ b/docs/zh/getting_started/README.md @@ -32,13 +32,13 @@ MAX_JOBS=8 uv pip install flash-attn --no-build-isolation CUDA 12.9 环境将 `cu130` 换成 `cu129`。已验证的 CUDA 12.9 [依赖 lock](../../../requirements/locks/README.md) 为可选复现方式。 -`einops`、`sgl-kernel` 以及训练、benchmark 和测试包均已是主依赖， +`einops`、`sglang-kernel` 以及训练、benchmark 和测试包均已是主依赖， 不再需要工作流专用 extra。 Sparse-vLLM 当前支持未量化 BF16 和 block-scaled FP8 格式的 Qwen3.5/Qwen3.6 checkpoint。 其 prefill causal Conv1D 和 decode Conv1D/GDN packing path 使用仓库本地 -Triton kernel，本身不调用 `sgl-kernel`。 +Triton kernel，本身不调用 `sglang-kernel`。 `flashinfer-cubin` 是可选加速 package： diff --git a/docs/zh/getting_started/reproducibility.md b/docs/zh/getting_started/reproducibility.md index 685e2891..26387106 100644 --- a/docs/zh/getting_started/reproducibility.md +++ b/docs/zh/getting_started/reproducibility.md @@ -9,10 +9,10 @@ README 包含当前安装命令。预期 baseline 为： - Python 3.10。 - 使用 `requirements/locks/canonical-cu129-py310.txt` 中冻结的完整 runtime 与 test 环境。 -- 带 CUDA 12.9 wheel 的 PyTorch 2.9.1，以及 Triton 3.5.1。 +- 带匹配 CUDA wheel 的 PyTorch 2.11.0，以及 Triton 3.6.0。 - `flashinfer-python==0.6.15.post1`，以及 CUDA 12.9 build 的 `flashinfer-jit-cache==0.6.15.post1`。 -- `sgl-kernel==0.3.21` 和 `einops>=0.8.2` 是 runtime 依赖。 +- `sglang-kernel==0.4.5` 和 `einops>=0.8.2` 是 runtime 依赖。 - 需要预编译 device binary 时，可从通用 FlashInfer wheel index 安装匹配的 `flashinfer-cubin`。 - `transformers==5.13.1`。 - 使用 `MAX_JOBS=8 pip install flash-attn --no-build-isolation` 安装 `flash-attn`。 diff --git a/pyproject.toml b/pyproject.toml index 7f97289c..3c111c58 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,7 +22,7 @@ dependencies = [ "nvidia-cutlass-dsl>=4.6,<5", "pillow", "einops", - "sgl-kernel>=0.3.21,<0.4", + "sglang-kernel>=0.4.5,<0.4.6", "tqdm", "loguru", "fastapi>=0.100", @@ -45,12 +45,12 @@ dependencies = [ [project.optional-dependencies] cu129 = [ - "torch>=2.9,<3", + "torch==2.11.0", "flashinfer-python[cu12]>=0.6.15,<0.7", "flashinfer-jit-cache>=0.6.15,<0.7", ] cu130 = [ - "torch>=2.9,<3", + "torch==2.11.0", "flashinfer-python[cu13]>=0.6.15,<0.7", "flashinfer-jit-cache>=0.6.15,<0.7", ] diff --git a/src/sparsevllm/kernels/external/sgl/fa3.py b/src/sparsevllm/kernels/external/sgl/fa3.py index db1df9c2..a513ab08 100644 --- a/src/sparsevllm/kernels/external/sgl/fa3.py +++ b/src/sparsevllm/kernels/external/sgl/fa3.py @@ -45,6 +45,8 @@ "pack_gqa", "sm_margin", "sinks", + "sparse_mask_fine", + "only_qv", ) _FWD_ARGUMENTS = _COMMON_FWD_ARGUMENTS + ("attention_chunk",) + _FWD_TRAILING_ARGUMENTS @@ -114,11 +116,11 @@ def sgl_fa3_support() -> tuple[bool, str]: argument_names = tuple(argument.name for argument in op._schema.arguments) except Exception as error: return False, ( - "sgl-kernel FA3 raw op failed to load: " + "sglang-kernel FA3 raw op failed to load: " f"{type(error).__name__}: {error}" ) if argument_names != _FWD_ARGUMENTS: - return False, f"unsupported sgl-kernel FA3 fwd schema: {argument_names}" + return False, f"unsupported sglang-kernel FA3 fwd schema: {argument_names}" return True, reason @@ -356,11 +358,11 @@ def run_varlen( ] args.append(0) args.extend( - (0.0, True, scheduler_metadata, split_count, None, 0, None) + (0.0, True, scheduler_metadata, split_count, None, 0, None, None, False) ) result: Sequence[torch.Tensor] = self._op(*args) if not result or result[0].data_ptr() != output.data_ptr(): - raise RuntimeError("sgl-kernel FA3 did not write to the supplied output") + raise RuntimeError("sglang-kernel FA3 did not write to the supplied output") return output def run_explicit_varlen( @@ -422,11 +424,11 @@ def run_explicit_varlen( ] args.append(0) args.extend( - (0.0, True, scheduler_metadata, self.num_splits, None, 0, None) + (0.0, True, scheduler_metadata, self.num_splits, None, 0, None, None, False) ) result: Sequence[torch.Tensor] = self._op(*args) if not result or result[0].data_ptr() != output.data_ptr(): - raise RuntimeError("sgl-kernel FA3 did not write to the supplied output") + raise RuntimeError("sglang-kernel FA3 did not write to the supplied output") return output def run_contiguous_explicit_varlen( @@ -473,10 +475,10 @@ def run_contiguous_explicit_varlen( -1, ] args.append(0) - args.extend((0.0, True, None, self.num_splits, None, 0, None)) + args.extend((0.0, True, None, self.num_splits, None, 0, None, None, False)) result: Sequence[torch.Tensor] = self._op(*args) if not result or result[0].data_ptr() != output.data_ptr(): - raise RuntimeError("sgl-kernel FA3 did not write to the supplied output") + raise RuntimeError("sglang-kernel FA3 did not write to the supplied output") return output diff --git a/src/sparsevllm/kernels/external/sgl/moe.py b/src/sparsevllm/kernels/external/sgl/moe.py index 5f710837..dfb030d1 100644 --- a/src/sparsevllm/kernels/external/sgl/moe.py +++ b/src/sparsevllm/kernels/external/sgl/moe.py @@ -19,12 +19,12 @@ def sgl_moe_alignment_support() -> tuple[bool, str]: alignment = importlib.import_module("sgl_kernel").moe_align_block_size except Exception as error: return False, ( - "sgl-kernel MoE alignment failed to load: " + "sglang-kernel MoE alignment failed to load: " f"{type(error).__name__}: {error}" ) return (True, reason) if callable(alignment) else ( False, - "sgl-kernel moe_align_block_size is not callable", + "sglang-kernel moe_align_block_size is not callable", ) @@ -69,8 +69,8 @@ def sgl_moe_align_block_size( ) from sgl_kernel import moe_align_block_size - # sgl-kernel 0.3.x iterates to num_experts - 1. The extra empty logical - # expert makes the complete [0, num_experts) range participate. + # The extra empty logical expert makes the complete [0, num_experts) + # range participate; hardware tests cover assignments to the final expert. moe_align_block_size( topk_ids, num_experts + 1, diff --git a/src/sparsevllm/kernels/external/sgl/support.py b/src/sparsevllm/kernels/external/sgl/support.py index ba1a18b2..e26ca8d8 100644 --- a/src/sparsevllm/kernels/external/sgl/support.py +++ b/src/sparsevllm/kernels/external/sgl/support.py @@ -5,8 +5,9 @@ import importlib.util import re -_MIN_VERSION = (0, 3, 21) -_MAX_VERSION = (0, 4, 0) +_MIN_VERSION = (0, 4, 5) +_MAX_VERSION = (0, 4, 6) +_DISTRIBUTION = "sglang-kernel" def sgl_kernel_support(feature: str) -> tuple[bool, str]: @@ -15,22 +16,22 @@ def sgl_kernel_support(feature: str) -> tuple[bool, str]: try: package_spec = importlib.util.find_spec("sgl_kernel") except (ImportError, ValueError) as error: - return False, f"sgl-kernel package discovery failed: {error}" + return False, f"{_DISTRIBUTION} package discovery failed: {error}" if package_spec is None: - return False, "sgl-kernel is not installed" + return False, f"{_DISTRIBUTION} is not installed" try: - version = importlib.metadata.version("sgl-kernel") + version = importlib.metadata.version(_DISTRIBUTION) except importlib.metadata.PackageNotFoundError: - return False, "sgl-kernel package metadata is unavailable" + return False, f"{_DISTRIBUTION} package metadata is unavailable" match = re.match(r"^(\d+)\.(\d+)\.(\d+)", version) parsed = tuple(map(int, match.groups())) if match else None if parsed is None or not _MIN_VERSION <= parsed < _MAX_VERSION: - return False, f"requires sgl-kernel>=0.3.21,<0.4, got {version}" + return False, f"requires {_DISTRIBUTION}>=0.4.5,<0.4.6, got {version}" try: importlib.import_module("sgl_kernel") except Exception as error: return False, ( - f"sgl-kernel {version} {feature} failed to load: " + f"{_DISTRIBUTION} {version} {feature} failed to load: " f"{type(error).__name__}: {error}" ) - return True, f"sgl-kernel {version} {feature} is available" + return True, f"{_DISTRIBUTION} {version} {feature} is available" diff --git a/tests/test_dependency_constraints.py b/tests/test_dependency_constraints.py index ffefdd58..1e3187e0 100644 --- a/tests/test_dependency_constraints.py +++ b/tests/test_dependency_constraints.py @@ -15,7 +15,8 @@ def test_runtime_compatibility_bounds_cover_canonical_lock(): assert "tilelang==0.1.9" in dependencies assert "apache-tvm-ffi==0.1.10" in dependencies assert "transformers>=5.13,<6" in dependencies - assert "sgl-kernel>=0.3.21,<0.4" in dependencies + assert "nvidia-cutlass-dsl>=4.6,<5" in dependencies + assert "sglang-kernel>=0.4.5,<0.4.6" in dependencies assert {"fire", "pillow", "einops", "tqdm", "loguru"} <= dependencies assert not any( dependency.startswith("torchvision") for dependency in dependencies @@ -50,12 +51,12 @@ def test_uv_routes_cuda_packages_to_explicit_indexes(): assert config["project"]["optional-dependencies"] == { "cu129": [ - "torch>=2.9,<3", + "torch==2.11.0", "flashinfer-python[cu12]>=0.6.15,<0.7", "flashinfer-jit-cache>=0.6.15,<0.7", ], "cu130": [ - "torch>=2.9,<3", + "torch==2.11.0", "flashinfer-python[cu13]>=0.6.15,<0.7", "flashinfer-jit-cache>=0.6.15,<0.7", ], diff --git a/tests/test_mla_attention_operator.py b/tests/test_mla_attention_operator.py index 42996b89..e55048d1 100644 --- a/tests/test_mla_attention_operator.py +++ b/tests/test_mla_attention_operator.py @@ -102,7 +102,7 @@ def test_mla_resolver_selects_h100_triton_provider() -> None: with ( patch( "sparsevllm.operators.mla_attention.sgl_fa3_support", - return_value=(False, "sgl-kernel is not installed"), + return_value=(False, "sglang-kernel is not installed"), ), patch( "sparsevllm.operators.mla_attention.allocate_mla_decode_workspace", @@ -119,8 +119,8 @@ def test_mla_resolver_selects_h100_triton_provider() -> None: assert isinstance(resolved.provider, MlaTritonProvider) assert resolved.rejected == ( - ("sgl_fa3_h100", "sgl-kernel is not installed"), - ("tilelang_score_sgl_fa3_h100", "sgl-kernel is not installed"), + ("sgl_fa3_h100", "sglang-kernel is not installed"), + ("tilelang_score_sgl_fa3_h100", "sglang-kernel is not installed"), ) allocate.assert_called_once_with( batch_size=8, diff --git a/tests/test_sgl_fa3.py b/tests/test_sgl_fa3.py index 30f28b28..adaf1eb3 100644 --- a/tests/test_sgl_fa3.py +++ b/tests/test_sgl_fa3.py @@ -15,32 +15,23 @@ def test_sgl_fa3_support_rejects_missing_package() -> None: with patch("importlib.util.find_spec", return_value=None): - assert sgl_fa3_support() == (False, "sgl-kernel is not installed") + assert sgl_fa3_support() == (False, "sglang-kernel is not installed") -def test_sgl_fa3_support_rejects_old_package() -> None: +@pytest.mark.parametrize("version", ["0.4.4", "0.4.6"]) +def test_sgl_fa3_support_rejects_outside_declared_range(version: str) -> None: with ( patch("importlib.util.find_spec", return_value=object()), - patch("importlib.metadata.version", return_value="0.3.20"), + patch("importlib.metadata.version", return_value=version), ): supported, reason = sgl_fa3_support() assert not supported - assert ">=0.3.21,<0.4" in reason + assert "sglang-kernel>=0.4.5,<0.4.6" in reason -def test_sgl_fa3_support_rejects_04_package() -> None: - with ( - patch("importlib.util.find_spec", return_value=object()), - patch("importlib.metadata.version", return_value="0.4.5"), - ): - supported, reason = sgl_fa3_support() - - assert not supported - assert ">=0.3.21,<0.4" in reason - - -def test_sgl_fa3_support_accepts_declared_package() -> None: +def test_sgl_fa3_support_accepts_declared_range() -> None: + version = "0.4.5" op = SimpleNamespace( _schema=SimpleNamespace( arguments=[ @@ -55,19 +46,19 @@ def test_sgl_fa3_support_accepts_declared_package() -> None: return_value=op, ), patch("importlib.util.find_spec", return_value=object()), - patch("importlib.metadata.version", return_value="0.3.21"), + patch("importlib.metadata.version", return_value=version), patch("importlib.import_module", return_value=object()), ): supported, reason = sgl_fa3_support() assert supported - assert "0.3.21" in reason + assert version in reason def test_sgl_fa3_support_rejects_binary_load_failure() -> None: with ( patch("importlib.util.find_spec", return_value=object()), - patch("importlib.metadata.version", return_value="0.3.21"), + patch("importlib.metadata.version", return_value="0.4.5"), patch( "importlib.import_module", side_effect=ImportError("undefined symbol: c10_cuda_check"), @@ -98,7 +89,7 @@ def test_sgl_fa3_support_rejects_missing_op_schema() -> None: @pytest.mark.skipif( not torch.cuda.is_available() or not sgl_fa3_support()[0], - reason="CUDA and a validated sgl-kernel are required", + reason="CUDA and a validated sglang-kernel are required", ) def test_sgl_fa3_decode_matches_torch_and_replays_cuda_graph() -> None: torch.manual_seed(20260807) @@ -238,7 +229,7 @@ def counted_scheduler_op(*args, **kwargs): @pytest.mark.skipif( not torch.cuda.is_available() or not sgl_fa3_support()[0], - reason="CUDA and a validated sgl-kernel are required", + reason="CUDA and a validated sglang-kernel are required", ) def test_sgl_fa3_varlen_latent_prefill_matches_causal_torch() -> None: torch.manual_seed(20260807) @@ -310,7 +301,7 @@ def test_sgl_fa3_varlen_latent_prefill_matches_causal_torch() -> None: @pytest.mark.skipif( not torch.cuda.is_available() or not sgl_fa3_support()[0], - reason="CUDA and a validated sgl-kernel are required", + reason="CUDA and a validated sglang-kernel are required", ) def test_sgl_fa3_varlen_explicit_prefill_matches_causal_torch() -> None: torch.manual_seed(20260807) diff --git a/tests/test_sgl_moe.py b/tests/test_sgl_moe.py index 825522ec..82386f19 100644 --- a/tests/test_sgl_moe.py +++ b/tests/test_sgl_moe.py @@ -18,38 +18,40 @@ def test_sgl_moe_support_rejects_missing_package() -> None: with patch("importlib.util.find_spec", return_value=None): assert sgl_moe_alignment_support() == ( False, - "sgl-kernel is not installed", + "sglang-kernel is not installed", ) -def test_sgl_moe_support_rejects_04_package() -> None: +@pytest.mark.parametrize("version", ["0.4.4", "0.4.6"]) +def test_sgl_moe_support_rejects_outside_declared_range(version: str) -> None: with ( patch("importlib.util.find_spec", return_value=object()), - patch("importlib.metadata.version", return_value="0.4.5"), + patch("importlib.metadata.version", return_value=version), ): supported, reason = sgl_moe_alignment_support() assert not supported - assert ">=0.3.21,<0.4" in reason + assert "sglang-kernel>=0.4.5,<0.4.6" in reason -def test_sgl_moe_support_accepts_declared_package() -> None: +def test_sgl_moe_support_accepts_declared_range() -> None: + version = "0.4.5" module = SimpleNamespace(moe_align_block_size=lambda *_args: None) with ( patch("importlib.util.find_spec", return_value=object()), - patch("importlib.metadata.version", return_value="0.3.21"), + patch("importlib.metadata.version", return_value=version), patch("importlib.import_module", return_value=module), ): supported, reason = sgl_moe_alignment_support() assert supported - assert "0.3.21" in reason + assert version in reason def test_sgl_moe_support_rejects_missing_alignment_api() -> None: with ( patch("importlib.util.find_spec", return_value=object()), - patch("importlib.metadata.version", return_value="0.3.21"), + patch("importlib.metadata.version", return_value="0.4.5"), patch("importlib.import_module", return_value=SimpleNamespace()), ): supported, reason = sgl_moe_alignment_support() @@ -60,7 +62,7 @@ def test_sgl_moe_support_rejects_missing_alignment_api() -> None: @pytest.mark.skipif( not torch.cuda.is_available() or not sgl_moe_alignment_support()[0], - reason="CUDA and a validated sgl-kernel are required", + reason="CUDA and a validated sglang-kernel are required", ) def test_sgl_moe_alignment_matches_grouped_assignments() -> None: torch.manual_seed(20260808) @@ -107,7 +109,7 @@ def test_sgl_moe_alignment_matches_grouped_assignments() -> None: @pytest.mark.skipif( not torch.cuda.is_available() or not sgl_moe_alignment_support()[0], - reason="CUDA and a validated sgl-kernel are required", + reason="CUDA and a validated sglang-kernel are required", ) @pytest.mark.parametrize(("local_start", "local_end"), [(0, 32), (32, 64)]) def test_sgl_moe_alignment_matches_ep_shard( @@ -170,7 +172,7 @@ def grouped_assignments(sorted_ids, expert_ids, count): @pytest.mark.skipif( not torch.cuda.is_available() or not sgl_moe_alignment_support()[0], - reason="CUDA and a validated sgl-kernel are required", + reason="CUDA and a validated sglang-kernel are required", ) @pytest.mark.parametrize("num_tokens", [1, 2, 4, 5, 64]) @pytest.mark.parametrize("all_remote", [False, True]) From 01ed2aad0565c562f020d3306b99410e2ab0fb10 Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Thu, 13 Aug 2026 21:49:50 +0800 Subject: [PATCH 44/48] fix: bound cuda graph context buckets --- src/sparsevllm/configs/cuda_graph.py | 2 +- tests/test_prefill_schedule_policy.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/sparsevllm/configs/cuda_graph.py b/src/sparsevllm/configs/cuda_graph.py index 3e99aa92..548b6cb7 100644 --- a/src/sparsevllm/configs/cuda_graph.py +++ b/src/sparsevllm/configs/cuda_graph.py @@ -134,7 +134,7 @@ def _default_decode_cuda_graph_context_sizes(max_model_len: int) -> list[int]: while size < max_model_len: sizes.append(size) size *= 2 - sizes.append(size) + sizes.append(max_model_len) return sorted(set(sizes)) diff --git a/tests/test_prefill_schedule_policy.py b/tests/test_prefill_schedule_policy.py index 16652274..621686d6 100644 --- a/tests/test_prefill_schedule_policy.py +++ b/tests/test_prefill_schedule_policy.py @@ -1321,7 +1321,7 @@ def test_decode_cuda_graph_auto_context_sizes_use_powers_of_two_from_1k(self): decode_cuda_graph=True, max_model_len=9000, ) - self.assertEqual(cfg.decode_cuda_graph_context_sizes, [1024, 2048, 4096, 8192, 16384]) + self.assertEqual(cfg.decode_cuda_graph_context_sizes, [1024, 2048, 4096, 8192, 9000]) def test_decode_cuda_graph_explicit_context_sizes_are_sorted(self): cfg = self.make_config( From 183ae5f5da4fa5ccf783961bd157b7495d8ffbf7 Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Thu, 13 Aug 2026 21:50:08 +0800 Subject: [PATCH 45/48] fix: correct sgl moe expert alignment --- src/sparsevllm/operators/moe.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/sparsevllm/operators/moe.py b/src/sparsevllm/operators/moe.py index 31ad8679..33605ff2 100644 --- a/src/sparsevllm/operators/moe.py +++ b/src/sparsevllm/operators/moe.py @@ -243,7 +243,7 @@ def _sgl_moe_align_block_size( return sgl_moe_align_block_size( topk_ids, block_size=block_size, - num_experts=num_local_experts + int(has_remote_experts), + num_experts=num_local_experts, ) From ef453047c815573611c26050e4ab1a6cfb70c8a5 Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Thu, 13 Aug 2026 21:50:17 +0800 Subject: [PATCH 46/48] feat: enable glm providers on h20 --- src/sparsevllm/operators/mla_attention.py | 20 +++++++++--- src/sparsevllm/operators/moe.py | 4 +-- tests/test_mla_attention_operator.py | 39 +++++++++++++++++------ tests/test_operator_providers.py | 15 +++++---- 4 files changed, 56 insertions(+), 22 deletions(-) diff --git a/src/sparsevllm/operators/mla_attention.py b/src/sparsevllm/operators/mla_attention.py index 4ad1c96d..5ce6fe47 100644 --- a/src/sparsevllm/operators/mla_attention.py +++ b/src/sparsevllm/operators/mla_attention.py @@ -38,6 +38,8 @@ _GLM_MLA_QK_HEAD_DIM = 256 _GLM_MLA_VALUE_HEAD_DIM = 256 _VALIDATED_H100_NAME = "NVIDIA H100 80GB HBM3" +_VALIDATED_H20_NAME = "NVIDIA H20" +_VALIDATED_SM90_NAMES = frozenset({_VALIDATED_H100_NAME, _VALIDATED_H20_NAME}) @dataclass(frozen=True, slots=True) @@ -107,9 +109,9 @@ def run( @MLA_ATTENTION_REGISTRY.register class MlaTritonProvider(MlaAttentionProvider): - """H100 provider bound once with caller-independent decode workspace.""" + """Validated SM90 provider with caller-independent decode workspace.""" - name = "triton_h100" + name = "triton_sm90" priority = 100 def __init__( @@ -162,9 +164,9 @@ def supports( f"requires CUDA SM90, got {caps.platform.name} " f"{caps.compute_capability}" ) - if caps.device_name != _VALIDATED_H100_NAME: + if caps.device_name not in _VALIDATED_SM90_NAMES: return SupportResult.no( - "requires validated NVIDIA H100 80GB HBM3 hardware, got " + "requires validated H100 80GB HBM3 or H20 hardware, got " f"{caps.device_name}" ) if not caps.supports_triton: @@ -205,6 +207,10 @@ def supports( return SupportResult.no( f"requires tensor parallel size 1, 2, or 4, got {spec.tp_size}" ) + if caps.device_name == _VALIDATED_H20_NAME and spec.tp_size not in {1, 2}: + return SupportResult.no( + f"H20 MLA currently requires tensor parallel size 1 or 2, got {spec.tp_size}" + ) return SupportResult.yes() def _validate_run_inputs( @@ -403,7 +409,7 @@ def run( class MlaSglFa3Provider(MlaTritonProvider): """SGL FA3 decode with the score-producing Triton path kept explicit.""" - name = "sgl_fa3_h100" + name = "sgl_fa3_sm90" priority = 200 supports_explicit_prefill = True @@ -651,6 +657,10 @@ def supports( base = MlaSglFa3Provider.supports(spec, caps) if not base.supported: return base + if caps.device_name != _VALIDATED_H100_NAME: + return SupportResult.no( + "requires H100-validated TileLang MLA schedules" + ) supported, reason = tilelang_mla_support() return SupportResult.yes(reason) if supported else SupportResult.no(reason) diff --git a/src/sparsevllm/operators/moe.py b/src/sparsevllm/operators/moe.py index 33605ff2..74eec8ba 100644 --- a/src/sparsevllm/operators/moe.py +++ b/src/sparsevllm/operators/moe.py @@ -257,8 +257,8 @@ class SglAlignedTritonGlmMoeProvider(MoeProvider): def supports(cls, spec: MoeOpSpec, caps: DeviceCaps) -> SupportResult: if caps.platform != PlatformEnum.CUDA or caps.compute_capability != (9, 0): return SupportResult.no("requires CUDA SM90") - if caps.device_name != "NVIDIA H100 80GB HBM3": - return SupportResult.no("requires profiled NVIDIA H100 80GB HBM3") + if caps.device_name not in {"NVIDIA H100 80GB HBM3", "NVIDIA H20"}: + return SupportResult.no("requires validated H100 80GB HBM3 or H20 hardware") expected = { (64, 64, 2048, 768, 4, 2, 1, "biased_sigmoid"), (65, 65, 2048, 768, 5, 2, 1, "biased_sigmoid"), diff --git a/tests/test_mla_attention_operator.py b/tests/test_mla_attention_operator.py index e55048d1..4672e38e 100644 --- a/tests/test_mla_attention_operator.py +++ b/tests/test_mla_attention_operator.py @@ -95,8 +95,14 @@ def test_mla_attention_scale_uses_qk_head_dimension() -> None: assert spec.softmax_scale != pytest.approx((512 + 64) ** -0.5) -def test_mla_resolver_selects_h100_triton_provider() -> None: - spec = _spec() +@pytest.mark.parametrize( + ("device_name", "tp_size"), + [("NVIDIA H100 80GB HBM3", 4), ("NVIDIA H20", 1), ("NVIDIA H20", 2)], +) +def test_mla_resolver_selects_sm90_triton_provider( + device_name: str, tp_size: int +) -> None: + spec = _spec(tp_size=tp_size) workspace = _cpu_workspace(batch_size=8, head_count=spec.local_q_heads) with ( @@ -111,7 +117,7 @@ def test_mla_resolver_selects_h100_triton_provider() -> None: ): resolved = OpResolver(MLA_ATTENTION_REGISTRY).resolve( spec, - _h100_caps(), + _h100_caps(device_name=device_name), op_spec=spec, device="cpu", max_batch_size=8, @@ -119,19 +125,29 @@ def test_mla_resolver_selects_h100_triton_provider() -> None: assert isinstance(resolved.provider, MlaTritonProvider) assert resolved.rejected == ( - ("sgl_fa3_h100", "sglang-kernel is not installed"), + ("sgl_fa3_sm90", "sglang-kernel is not installed"), ("tilelang_score_sgl_fa3_h100", "sglang-kernel is not installed"), ) allocate.assert_called_once_with( batch_size=8, - head_count=5, + head_count=spec.local_q_heads, device=torch.device("cpu"), config=GLM_MLA_MAX_WORKSPACE_CONFIG, ) -def test_mla_resolver_prefers_sgl_fa3_when_available() -> None: - spec = _spec() +@pytest.mark.parametrize( + ("device_name", "tp_size", "tilelang_reason"), + [ + ("NVIDIA H100 80GB HBM3", 4, "tilelang unavailable"), + ("NVIDIA H20", 1, "requires H100-validated TileLang MLA schedules"), + ("NVIDIA H20", 2, "requires H100-validated TileLang MLA schedules"), + ], +) +def test_mla_resolver_prefers_sgl_fa3_on_supported_sm90( + device_name: str, tp_size: int, tilelang_reason: str +) -> None: + spec = _spec(tp_size=tp_size) workspace = _cpu_workspace(batch_size=8, head_count=spec.local_q_heads) with ( @@ -151,7 +167,7 @@ def test_mla_resolver_prefers_sgl_fa3_when_available() -> None: ): resolved = OpResolver(MLA_ATTENTION_REGISTRY).resolve( spec, - _h100_caps(), + _h100_caps(device_name=device_name), op_spec=spec, device="cpu", max_batch_size=8, @@ -161,7 +177,7 @@ def test_mla_resolver_prefers_sgl_fa3_when_available() -> None: assert resolved.rejected == ( ( "tilelang_score_sgl_fa3_h100", - "tilelang unavailable", + tilelang_reason, ), ) @@ -204,6 +220,11 @@ def test_mla_resolver_accepts_cuda_graph_after_capture_gate() -> None: ({}, {"platform": PlatformEnum.CPU}, "requires CUDA SM90"), ({}, {"compute_capability": (8, 0)}, "requires CUDA SM90"), ({}, {"device_name": "NVIDIA H100 PCIe"}, "validated"), + ( + {"tp_size": 4}, + {"device_name": "NVIDIA H20"}, + "H20 MLA currently requires tensor parallel size 1 or 2", + ), ({}, {"supports_triton": False}, "does not support Triton"), ({}, {"supports_bfloat16": False}, "does not support BF16"), ( diff --git a/tests/test_operator_providers.py b/tests/test_operator_providers.py index 58a18b9e..29fced49 100644 --- a/tests/test_operator_providers.py +++ b/tests/test_operator_providers.py @@ -567,7 +567,8 @@ def test_hopper_fused_moe_uses_profiled_tp_ep_shape(): assert resolved.provider.name == "triton_hopper_fused" -def test_glm_tp2_moe_uses_sgl_aligned_provider(): +@pytest.mark.parametrize("device_name", ["NVIDIA H100 80GB HBM3", "NVIDIA H20"]) +def test_glm_tp2_moe_uses_sgl_aligned_provider(device_name): spec = _moe_spec( activation_dtype=torch.bfloat16, weight_dtype=torch.bfloat16, @@ -591,14 +592,15 @@ def test_glm_tp2_moe_uses_sgl_aligned_provider(): _cuda_caps( (9, 0), native_fp8=False, - device_name="NVIDIA H100 80GB HBM3", + device_name=device_name, ), ) assert resolved.provider.name == "sgl_aligned_triton_glm" -def test_glm_tp2_fused_shared_decode_uses_sgl_aligned_provider(): +@pytest.mark.parametrize("device_name", ["NVIDIA H100 80GB HBM3", "NVIDIA H20"]) +def test_glm_tp2_fused_shared_decode_uses_sgl_aligned_provider(device_name): spec = _moe_spec( activation_dtype=torch.bfloat16, weight_dtype=torch.bfloat16, @@ -622,7 +624,7 @@ def test_glm_tp2_fused_shared_decode_uses_sgl_aligned_provider(): _cuda_caps( (9, 0), native_fp8=False, - device_name="NVIDIA H100 80GB HBM3", + device_name=device_name, ), ) @@ -654,7 +656,8 @@ def test_glm_packed_shared_expert_profile(overrides, expected): assert use_packed_shared_experts(**values) is expected -def test_glm_tp2_ep2_moe_uses_sgl_aligned_provider(): +@pytest.mark.parametrize("device_name", ["NVIDIA H100 80GB HBM3", "NVIDIA H20"]) +def test_glm_tp2_ep2_moe_uses_sgl_aligned_provider(device_name): spec = _moe_spec( activation_dtype=torch.bfloat16, weight_dtype=torch.bfloat16, @@ -678,7 +681,7 @@ def test_glm_tp2_ep2_moe_uses_sgl_aligned_provider(): _cuda_caps( (9, 0), native_fp8=False, - device_name="NVIDIA H100 80GB HBM3", + device_name=device_name, ), ) From 4d229fbbadbf2253b56faddddaa06bbdb90a85fb Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Thu, 13 Aug 2026 21:52:04 +0800 Subject: [PATCH 47/48] test: initialize graph runner counter --- tests/test_prefill_schedule_policy.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_prefill_schedule_policy.py b/tests/test_prefill_schedule_policy.py index 621686d6..aa88e4ee 100644 --- a/tests/test_prefill_schedule_policy.py +++ b/tests/test_prefill_schedule_policy.py @@ -1350,6 +1350,7 @@ def make_runner(self, method="quest", cache_manager=None): runner.recurrent_state_manager = None runner.max_context_len_override = None runner._graphs = {} + runner.eager_static_count = 0 runner.capture_sizes = [1, 2, 4, 8, 16] runner.context_sizes = [1024, 2048, 4096, 8192, 16384, 32768, 65536, 131072] return runner From 97a3051a498aa30762219b30137aaa5c12afbe35 Mon Sep 17 00:00:00 2001 From: QuanshengGu Date: Thu, 13 Aug 2026 17:27:56 +0800 Subject: [PATCH 48/48] fix: restore cpu test collection --- .../kernels/tilelang/mla/__init__.py | 6 +----- tests/test_prefill_schedule_policy.py | 1 + tests/test_snapkv_cache_budget.py | 1 + tests/test_tilelang_mla_operator.py | 18 ++++++++++++++++++ 4 files changed, 21 insertions(+), 5 deletions(-) diff --git a/src/sparsevllm/kernels/tilelang/mla/__init__.py b/src/sparsevllm/kernels/tilelang/mla/__init__.py index b6120d30..8d349656 100644 --- a/src/sparsevllm/kernels/tilelang/mla/__init__.py +++ b/src/sparsevllm/kernels/tilelang/mla/__init__.py @@ -1,5 +1 @@ -"""GLM MLA TileLang kernels.""" - -from .decode import build_glm_mla_decode_kernel, pad_glm_q_kernel - -__all__ = ["build_glm_mla_decode_kernel", "pad_glm_q_kernel"] +"""GLM MLA TileLang kernels loaded lazily by the runtime adapter.""" diff --git a/tests/test_prefill_schedule_policy.py b/tests/test_prefill_schedule_policy.py index aa88e4ee..4bc8ef04 100644 --- a/tests/test_prefill_schedule_policy.py +++ b/tests/test_prefill_schedule_policy.py @@ -1353,6 +1353,7 @@ def make_runner(self, method="quest", cache_manager=None): runner.eager_static_count = 0 runner.capture_sizes = [1, 2, 4, 8, 16] runner.context_sizes = [1024, 2048, 4096, 8192, 16384, 32768, 65536, 131072] + runner.eager_static_count = 0 return runner def make_seq(self, *, prompt_len=100, max_tokens=900, num_tokens=101): diff --git a/tests/test_snapkv_cache_budget.py b/tests/test_snapkv_cache_budget.py index b8fe1b8e..6474a008 100644 --- a/tests/test_snapkv_cache_budget.py +++ b/tests/test_snapkv_cache_budget.py @@ -39,6 +39,7 @@ def _manager_config(*, method: str, compression_interval: int = 1): torch_dtype=torch.float32, ), runtime_layout=RuntimeLayout.dense(2), + attention_cache_layout="explicit_kv", max_model_len=5, max_num_batched_tokens=10, max_num_seqs_in_gpu=3, diff --git a/tests/test_tilelang_mla_operator.py b/tests/test_tilelang_mla_operator.py index 8dbd059c..29d3870f 100644 --- a/tests/test_tilelang_mla_operator.py +++ b/tests/test_tilelang_mla_operator.py @@ -1,5 +1,7 @@ from __future__ import annotations +import subprocess +import sys from importlib import metadata from unittest.mock import Mock, patch @@ -96,6 +98,22 @@ def version(package: str) -> str: ) +def test_tilelang_runtime_import_does_not_require_tilelang() -> None: + code = """ +import sys + +class RejectTileLang: + def find_spec(self, fullname, path=None, target=None): + if fullname == "tilelang" or fullname.startswith("tilelang."): + raise ModuleNotFoundError("blocked optional tilelang import") + +sys.meta_path.insert(0, RejectTileLang()) +from sparsevllm.kernels.tilelang.mla.runtime import tilelang_mla_support +assert "tilelang" not in sys.modules +""" + subprocess.run([sys.executable, "-c", code], check=True) + + @pytest.mark.parametrize( ("version", "supported"), [("0.1.8", False), ("0.1.9", True), ("0.2.0+cu130", False)],