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/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..61b5f453
--- /dev/null
+++ b/.agents/skills/optimize-sparsevllm-kernel/references/operator-integration.md
@@ -0,0 +1,54 @@
+# 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/kernels/triton/`.
+- Keep repository-owned TileLang kernels under
+ `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,
+ 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..a0334855
--- /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/kernels/tilelang/` 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..c8f1f996
--- /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/kernels/triton/`.
+
+## 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/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/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 39f8b2f0..2b87ad4b 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 | ✅ |
@@ -93,8 +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. Install the package from the
-repository root using the runtime versions pinned in `pyproject.toml`.
+Sparse-vLLM requires Python 3.10 or newer. Default dependencies are declared in
+`pyproject.toml`.
### Conda
@@ -102,14 +103,10 @@ 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 .
+CUDA_VERSION=cu130
+python -m pip config --site set global.extra-index-url \
+ "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
@@ -125,44 +122,24 @@ 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 -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
```
-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:
+Use `cu129` instead of `cu130` for CUDA 12.9. The validated CUDA 12.9
+[dependency lock](requirements/locks/README.md) is optional.
-```bash
-pip install -e ".[prefix-offload]"
-```
+`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; `sglang-kernel` and a local CUDA extension
-build are not required.
+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/README_zh.md b/README_zh.md
index 021af344..1169af80 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,10 @@ 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 .
+CUDA_VERSION=cu130
+python -m pip config --site set global.extra-index-url \
+ "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
@@ -113,37 +109,16 @@ 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 .
+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
```
-以上显式索引分别用于安装 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]"
-```
+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/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/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/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/en/features/supported-models.md b/docs/en/features/supported-models.md
index 8ff710a8..fd374fe3 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 / 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 | ✅ |
@@ -25,18 +26,29 @@ 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 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
| Model | Vanilla | StreamingLLM | SnapKV | H2O | PyramidKV | OmniKV | QuEST | R-KV | SkipKV | DeltaKV |
@@ -46,6 +58,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 | ✅⁵ | ✅⁵ | ✅⁵ | Experimental⁴⁵ | — | ✅⁵ | — | ✅⁵ | — | — |
| Llama 3 / 3.1 | ✅ | ✅ | ✅ | Experimental⁴ | ✅ | ✅ | ✅ | ✅ | Selected checkpoint¹ | Compressor required² |
| MiniMax M2.7 | ✅ | ✅ | ✅ | Experimental⁴ | ✅ | ✅ | ✅ | ✅ | — | — |
@@ -63,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/en/getting_started/README.md b/docs/en/getting_started/README.md
index d279147b..b0b12550 100644
--- a/docs/en/getting_started/README.md
+++ b/docs/en/getting_started/README.md
@@ -10,14 +10,10 @@ 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 .
+CUDA_VERSION=cu130
+python -m pip config --site set global.extra-index-url \
+ "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
@@ -25,58 +21,32 @@ MAX_JOBS=8 pip install flash-attn --no-build-isolation
## Install with uv
-The project uses the CUDA 13.0 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 -e ".[cu130]"
# 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:
+Use `cu129` instead of `cu130` for CUDA 12.9. The validated CUDA 12.9
+[dependency lock](../../../requirements/locks/README.md) is optional.
-```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`, `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. They do not require `sglang-kernel` or a repository CUDA-extension
-build.
+kernels and do not call `sglang-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" \
- --index-url https://flashinfer.ai/whl/cu130
-
-# 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..13504757 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.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`.
+- `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[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/design/README.md b/docs/zh/design/README.md
index ac2c815a..b8df59db 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/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/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..0ec711cd
--- /dev/null
+++ b/docs/zh/design/glm-4.7-flash-support-plan.md
@@ -0,0 +1,131 @@
+# GLM-4.7-Flash 运行时支持设计
+
+本文说明 `glm4_moe_lite` 的稳定运行时契约、组件所有权和支持边界。模型与
+稀疏方法的汇总矩阵见[支持的模型](../features/supported-models.md)。
+
+## 设计边界
+
+GLM-4.7-Flash 使用 BF16 latent MLA。每层、每个 token 的持久化 attention
+cache 由 512 维 latent 和 64 维 RoPE key 组成，不持久化展开后的多头 K/V。
+
+实现遵循以下所有权边界：
+
+- 模型层负责 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 方法分支。
+
+Attention compute view 使用公共 metadata 与 tagged payload。显式 K/V 和
+MLA latent payload 是不同类型；消费者收到错误 payload 时必须 fail fast，
+不得通过互斥 optional tensor 或 metadata 字典绕过类型契约。
+
+## Attention 数据路径
+
+Prefill 按以下顺序执行：
+
+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 在该次调用结束后释放。
+
+因此，多 chunk prefill 的后续 chunk 必须能看到此前所有可见 token，而不是只
+看到当前 chunk。
+
+Decode 使用 absorbed query 直接读取 latent cache。Provider 拥有 Triton
+workspace、调度和数值 kernel；输出经过 value projection 重建到模型 hidden
+维度。静态 batch 中的 padding row 不得读写 `slot=-1`。
+
+## Cache 与前缀复用
+
+MLA storage 的显存容量按 `512 + 64` 个 BF16 value/token/layer 核算。所有
+allocation、free、reuse、eviction、slot copy 和 prefix replay 都必须同时处理
+latent 与 RoPE cache，且显存统计必须覆盖两个物理张量。
+
+Prefix Cache 的模式由稀疏方法决定：
+
+- vanilla 和 OmniKV 使用 radix prefix cache。
+- StreamingLLM、SnapKV、H2O 和 R-KV 使用 chain prefix cache。
+
+Prefix hit 只有在请求实际复用了 token、cache 状态和方法特定 metadata 时才算
+成功；仅完成请求不能证明 prefix 路径生效。
+
+## MoE 与并行布局
+
+模型第 0 层为 Dense，后续层为 routed MoE，并包含 shared expert。GLM 复用
+Qwen3-MoE 的 packed-expert 物理执行，但保留自己的 router、模型 topology 和
+checkpoint loader。
+
+已支持的 `(TP, EP)` 布局为：
+
+| 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 |
+
+联合 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` 整除。
+
+## Operator 与平台边界
+
+当前 MLA provider 支持 NVIDIA H100 80GB HBM3、BF16、SM90，以及 TP 1、2、4。
+Provider 在模型初始化时解析并绑定；kernel 执行失败必须直接暴露，不能在
+forward 中静默切换到 Torch 或其他 backend。
+
+Vendor kernel 固定来源为 LightLLM commit
+`65c174ee95ac6a6fd36b18b63d0b33d97e76b770`。本地 vendor 目录保留
+Apache-2.0 license、来源映射和修改说明。模型和 CacheManager 不直接依赖
+LightLLM Python runtime。
+
+## Serving 边界
+
+Chat Completions 与 Responses API 都使用 Transformers response parser。本地
+只提供 GLM 缺失的声明式 response template。Terminal EOS、stop boundary、
+流式增量和 raw parser text 的映射由通用 dispatcher/detokenizer 处理，不在
+GLM parser 中加入模型特判。
+
+## 不支持的组合
+
+以下组合必须在配置或模型构造前明确拒绝，不能落入默认实现：
+
+- H100 以外的 GPU、非 BF16 checkpoint 和量化权重。
+- `DP>1`、上述矩阵以外的 TP/EP 布局。
+- MTP/speculative decoding；loader 只精确跳过 checkpoint 中的 MTP 层。
+- Prefix offload。
+- PyramidKV、QuEST、SkipKV、DeltaKV，以及多个稀疏方法叠加。
+- 128K/202K 长上下文容量或吞吐支持声明。
+
+## 验证门禁
+
+支持声明至少需要覆盖以下可复现门禁：
+
+- 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/docs/zh/features/supported-models.md b/docs/zh/features/supported-models.md
index 44331922..af43594a 100644
--- a/docs/zh/features/supported-models.md
+++ b/docs/zh/features/supported-models.md
@@ -13,22 +13,33 @@
| 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 / 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 在 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 层。
+
## 稀疏方法支持
| 模型 | Vanilla | StreamingLLM | SnapKV | H2O | PyramidKV | OmniKV | QuEST | R-KV | SkipKV | DeltaKV |
@@ -38,6 +49,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 | ✅ | ✅ | ✅ | 实验性⁴ | ✅ | ✅ | ✅ | ✅ | — | — |
@@ -54,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/docs/zh/getting_started/README.md b/docs/zh/getting_started/README.md
index fe17d7f6..bc4f76c9 100644
--- a/docs/zh/getting_started/README.md
+++ b/docs/zh/getting_started/README.md
@@ -8,14 +8,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 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 .
+CUDA_VERSION=cu130
+python -m pip config --site set global.extra-index-url \
+ "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
@@ -23,50 +19,30 @@ MAX_JOBS=8 pip install flash-attn --no-build-isolation
## 使用 uv 安装
-项目使用 CUDA 13.0 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 -e ".[cu130]"
# Optional
MAX_JOBS=8 uv pip install flash-attn --no-build-isolation
```
-对于 Qwen3.5/Qwen3.6 mixed-attention inference，还需安装可选 Python 依赖：
+CUDA 12.9 环境将 `cu130` 换成 `cu129`。已验证的 CUDA 12.9
+[依赖 lock](../../../requirements/locks/README.md) 为可选复现方式。
-```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`、`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，不需要 `sglang-kernel` 或编译仓库 CUDA extension。
+其 prefill causal Conv1D 和 decode Conv1D/GDN packing path 使用仓库本地
+Triton kernel，本身不调用 `sglang-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" \
- --index-url https://flashinfer.ai/whl/cu130
-
-# 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..26387106 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 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`。
+- `sglang-kernel==0.4.5` 和 `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..3c111c58 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -14,40 +14,45 @@ authors = [
]
dependencies = [
"fire",
- "torch==2.11.0",
- "transformers==5.13.1",
- "triton==3.6.0",
- "torchvision==0.26.0",
+ "transformers>=5.13,<6",
+ "socksio>=1,<2",
+ "triton>=3.5,<4",
+ "tilelang==0.1.9",
+ "apache-tvm-ffi==0.1.10",
+ "nvidia-cutlass-dsl>=4.6,<5",
"pillow",
- "datasets",
- "matplotlib",
- "seaborn",
- "loguru",
- "ansible",
- "bitsandbytes",
+ "einops",
+ "sglang-kernel>=0.4.5,<0.4.6",
"tqdm",
+ "loguru",
"fastapi>=0.100",
"uvicorn[standard]",
"pydantic>=2",
- "pytest",
+ "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.optional-dependencies]
+cu129 = [
+ "torch==2.11.0",
+ "flashinfer-python[cu12]>=0.6.15,<0.7",
+ "flashinfer-jit-cache>=0.6.15,<0.7",
+]
+cu130 = [
+ "torch==2.11.0",
+ "flashinfer-python[cu13]>=0.6.15,<0.7",
+ "flashinfer-jit-cache>=0.6.15,<0.7",
]
[project.scripts]
@@ -61,3 +66,41 @@ include = ["sparsevllm*"]
[tool.pytest.ini_options]
testpaths = ["tests"]
+
+[tool.uv]
+conflicts = [
+ [
+ { extra = "cu129" },
+ { extra = "cu130" },
+ ],
+]
+
+[tool.uv.sources]
+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"
+explicit = true
diff --git a/requirements/locks/README.md b/requirements/locks/README.md
new file mode 100644
index 00000000..768f842d
--- /dev/null
+++ b/requirements/locks/README.md
@@ -0,0 +1,21 @@
+# Dependency locks
+
+`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:
+
+```bash
+uv venv --python 3.10
+uv pip install -r requirements/locks/canonical-cu129-py310.txt
+uv pip install --no-deps -e .
+```
+
+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/requirements/locks/canonical-cu129-py310.txt b/requirements/locks/canonical-cu129-py310.txt
new file mode 100644
index 00000000..57e372f6
--- /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 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/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/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/configs/cuda_graph.py b/src/sparsevllm/configs/cuda_graph.py
index 537e0c61..548b6cb7 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)
@@ -92,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/src/sparsevllm/configs/model.py b/src/sparsevllm/configs/model.py
index cf4a598d..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 "
@@ -132,6 +135,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)
+ config.attention_cache_layout = model_spec.attention_cache_layout
validate_checkpoint(
model_type,
outer_config=config.outer_hf_config,
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/configs/runtime.py b/src/sparsevllm/configs/runtime.py
index 94b3895a..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
@@ -74,6 +75,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/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/debug/tiny_random.py b/src/sparsevllm/debug/tiny_random.py
index 8a9fba23..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))
@@ -116,11 +121,21 @@ 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:
+ 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}."
@@ -179,7 +194,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/distributed/parallel_context.py b/src/sparsevllm/distributed/parallel_context.py
index eb41ae37..3afc5a44 100644
--- a/src/sparsevllm/distributed/parallel_context.py
+++ b/src/sparsevllm/distributed/parallel_context.py
@@ -1,21 +1,17 @@
from __future__ import annotations
-from dataclasses import dataclass, field, replace
+from dataclasses import dataclass
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
-
-
@dataclass(frozen=True)
class ParallelGroup:
process_group: dist.ProcessGroup | None
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):
@@ -95,10 +91,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(
@@ -248,25 +241,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
return _PARALLEL_CONTEXT
diff --git a/src/sparsevllm/engine/cache_manager/__init__.py b/src/sparsevllm/engine/cache_manager/__init__.py
index 79a1a02b..b80dc107 100644
--- a/src/sparsevllm/engine/cache_manager/__init__.py
+++ b/src/sparsevllm/engine/cache_manager/__init__.py
@@ -1,11 +1,33 @@
from __future__ import annotations
-from .base import CacheManager, DecodeComputeView, LayerBatchStates, PrefillComputeView, SparseSelection
+from .base import (
+ AttentionCacheWrite,
+ AttentionKeyComputeView,
+ AttentionPayload,
+ AttentionViewMeta,
+ CacheManager,
+ DecodeComputeView,
+ ExplicitKVPayload,
+ ExplicitKVWrite,
+ LayerBatchStates,
+ MlaLatentPayload,
+ MlaLatentWrite,
+ PrefillComputeView,
+ SparseSelection,
+)
__all__ = [
+ "AttentionCacheWrite",
+ "AttentionKeyComputeView",
+ "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 d4fffd6c..65bcab09 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
@@ -18,8 +18,9 @@
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
@@ -141,34 +142,83 @@ 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 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."""
+
+ 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."""
+
+ meta: AttentionViewMeta
+ payload: AttentionPayload
+
+
+@dataclass(frozen=True)
+class AttentionKeyComputeView:
+ """Physical payload plus the slots whose actual attention keys are needed."""
- 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
+ payload: AttentionPayload
+
+
+AttentionKeyMaterializer = Callable[[AttentionKeyComputeView], torch.Tensor]
class CacheManager(ABC):
@@ -195,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)
@@ -226,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,
@@ -417,7 +464,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
@@ -508,6 +555,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,
@@ -576,6 +628,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,
@@ -598,7 +667,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,
@@ -620,6 +692,133 @@ 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 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,
@@ -640,6 +839,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
@@ -665,7 +893,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,
@@ -675,14 +903,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=payload,
)
def collect_prefill_attention_score(
@@ -769,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.
@@ -913,21 +1167,24 @@ 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(
- 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=payload,
)
def get_decode_block_seq(self, layer_idx: int, default: int) -> int:
@@ -1264,7 +1521,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())
@@ -1285,7 +1543,36 @@ 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 = 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:
+ 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, ...]:
@@ -1329,8 +1616,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)
@@ -1355,8 +1648,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/deltakv_base.py b/src/sparsevllm/engine/cache_manager/deltakv_base.py
index 3d33acac..a270c715 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:
@@ -2742,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,
@@ -2957,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)))
@@ -2983,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)))
@@ -3326,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 895e8c7c..35b3042d 100644
--- a/src/sparsevllm/engine/cache_manager/deltakv_less_memory.py
+++ b/src/sparsevllm/engine/cache_manager/deltakv_less_memory.py
@@ -2,17 +2,18 @@
import math
import os
+from dataclasses import replace
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
@@ -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
@@ -2808,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)))
@@ -2905,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,
@@ -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):
@@ -3338,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],
@@ -3402,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 4f199069..7089db33 100644
--- a/src/sparsevllm/engine/cache_manager/h2o.py
+++ b/src/sparsevllm/engine/cache_manager/h2o.py
@@ -7,12 +7,13 @@
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
-from .base import PrefillComputeView
+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])
@@ -754,6 +757,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 +775,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,
@@ -1117,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(
@@ -1179,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
@@ -1202,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
@@ -1221,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 ac249dcd..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
@@ -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))
@@ -264,10 +327,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 +341,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 +354,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)
@@ -371,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),
@@ -425,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:
@@ -623,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),
@@ -696,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 13ed9325..39bb73cc 100644
--- a/src/sparsevllm/engine/cache_manager/snapkv.py
+++ b/src/sparsevllm/engine/cache_manager/snapkv.py
@@ -15,13 +15,21 @@
)
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
-from .base import CacheManager, LayerBatchStates, PrefillComputeView, SparseSelection
+from .base import (
+ AttentionCacheWrite,
+ CacheManager,
+ ExplicitKVPayload,
+ LayerBatchStates,
+ PrefillComputeView,
+ SparseSelection,
+)
from .raw_kv_offload import RawKVOffloadBuffer
+from .storage import ExplicitKVStorage, create_attention_cache_storage
_INT32_BYTES = 4
@@ -87,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
@@ -375,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)
@@ -398,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 (
@@ -969,8 +1071,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 +1087,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 +1117,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 +1128,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 +1145,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
@@ -2506,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(
@@ -2694,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 a4693903..74a26780 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_decode_cuda_graph_slot_mappings()
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..2e564758
--- /dev/null
+++ b/src/sparsevllm/engine/cache_manager/storage/__init__.py
@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from .base import AttentionCacheStorage, CacheLayout
+from .explicit_kv import ExplicitKVStorage
+
+if TYPE_CHECKING:
+ from .mla_latent import MlaLatentStorage
+
+
+def create_attention_cache_storage(
+ config: Any,
+ *,
+ num_kv_heads: int,
+ head_dim: int,
+) -> AttentionCacheStorage:
+ configured_layout = config.attention_cache_layout
+ 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:
+ 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),
+ dtype=dtype,
+ )
+ 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",
+ "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..784d9e37
--- /dev/null
+++ b/src/sparsevllm/engine/cache_manager/storage/base.py
@@ -0,0 +1,57 @@
+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 validate_slot_mappings(
+ self,
+ slot_mappings: tuple[torch.Tensor, ...],
+ ) -> None: ...
+
+ def store(
+ self,
+ layer_idx: int,
+ slot_mapping: torch.Tensor,
+ 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
new file mode 100644
index 00000000..83874094
--- /dev/null
+++ b/src/sparsevllm/engine/cache_manager/storage/explicit_kv.py
@@ -0,0 +1,205 @@
+from __future__ import annotations
+
+import torch
+
+from sparsevllm.kernels.triton.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 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,
+ 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 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
new file mode 100644
index 00000000..6c378849
--- /dev/null
+++ b/src/sparsevllm/engine/cache_manager/storage/mla_latent.py
@@ -0,0 +1,238 @@
+from __future__ import annotations
+
+import torch
+
+from sparsevllm.kernels.triton.mla.copy_latent import (
+ copy_latent_to_cache,
+ validate_copy_slot_mapping,
+)
+from sparsevllm.kernels.triton.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_store_calls_remaining: dict[
+ tuple[str, int, int], int
+ ] = {}
+
+ 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_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:
+ 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]),
+ )
+
+ 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,
+ 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)
+ remaining = self._validated_store_calls_remaining.get(
+ slot_mapping_key,
+ 0,
+ )
+ use_prevalidated_mapping = 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:
+ 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..6ac0e70d 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
@@ -87,7 +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,
capture_sizes: list[int],
context_sizes: list[int] | tuple[int, ...] | str | int | None = None,
graph_pool=None,
@@ -99,7 +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.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):
@@ -111,6 +110,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)
@@ -206,12 +209,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(
@@ -421,6 +421,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 +474,7 @@ def _capture(
if sparse_keepalive is not None:
keepalive.extend(sparse_keepalive())
state.keepalive = keepalive
+ self.capture_count += 1
return state
def run(
@@ -493,6 +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 += 1
return self.run_eager_static(seqs), None
graph_batch_size = self._select_graph_batch_size(real_batch_size)
@@ -515,6 +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 += 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 +530,7 @@ def run(
self._restore_sparse_state_refs(state)
with profiler.record("decode_cuda_graph_replay"):
state.graph.replay()
+ 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
@@ -530,6 +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 += 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 7caf2f42..c688334f 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
@@ -44,6 +45,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:
@@ -60,12 +66,16 @@
Qwen35MoeForCausalLM = None
-def _create_model(hf_config, model_spec: ModelSpec):
+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)
+ 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_"
@@ -145,7 +155,6 @@ def __init__(
self.parallel_context = init_parallel_context(
topology=config.parallel_topology,
)
-
# 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.
@@ -160,8 +169,22 @@ def __init__(
"decode_cuda_graph",
bool(getattr(config, "decode_cuda_graph", False)),
)
-
- self.model = _create_model(hf_config, config.model_spec)
+ decode_static_capture_sizes = _resolve_decode_cuda_graph_capture_sizes(
+ config.decode_cuda_graph_capture_sizes,
+ config.max_decoding_seqs,
+ )
+ 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
@@ -261,10 +284,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,
@@ -278,7 +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,
capture_sizes=decode_static_capture_sizes,
context_sizes=decode_static_context_sizes,
graph_pool=self.cuda_graph_pool,
@@ -312,6 +330,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(
@@ -756,6 +778,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)
@@ -792,10 +825,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")
@@ -806,12 +902,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)
@@ -834,6 +932,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),
@@ -895,18 +995,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()):
@@ -1159,10 +1270,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:
@@ -1241,6 +1361,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/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/engine/sparse_controller.py b/src/sparsevllm/engine/sparse_controller.py
index 417be53f..1f9d7214 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.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
@@ -13,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)
@@ -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/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/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 7bad90b2..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
@@ -58,8 +59,11 @@ 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
+ 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)
@@ -582,6 +586,14 @@ def _admit(self, item: _QueuedRequest, active: dict[int, _ActiveRequest]):
if isinstance(item.prompt, list)
else self.engine.tokenizer.encode(item.prompt)
)
+ 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
+ ),
+ )
active[seq_id] = _ActiveRequest(
index=item.index,
loop=item.loop,
@@ -593,6 +605,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 +663,46 @@ 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,
+ visible_text_len: int,
+ ) -> 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
+ ):
+ 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,
+ )
+
+ if not raw_text.startswith(parser_raw_text):
+ raise RuntimeError(
+ "Response-parser raw text is not a prefix of detokenized "
+ f"raw text: parser={parser_raw_text!r} raw={raw_text!r}."
+ )
+ 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 = {
seq_id: (token_logprobs, top_logprobs)
@@ -670,8 +726,7 @@ def _publish_token_deltas(self, active: dict[int, _ActiveRequest]):
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)
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]
@@ -680,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"
@@ -693,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(
@@ -704,7 +798,7 @@ def _publish_token_deltas(self, active: dict[int, _ActiveRequest]):
"type": "final",
"index": request.index,
"text": visible_text,
- "raw_text": final.raw_text,
+ "raw_text": final_parser_raw_text,
"text_delta": visible_text[request.emitted_text_len:],
"finish_reason": "stop",
"prompt_tokens": (
@@ -748,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,
@@ -977,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"
@@ -997,9 +1118,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 +1133,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/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 " 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 _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]:
+ 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, (
+ "sglang-kernel FA3 raw op failed to load: "
+ f"{type(error).__name__}: {error}"
+ )
+ if argument_names != _FWD_ARGUMENTS:
+ return False, f"unsupported sglang-kernel FA3 fwd schema: {argument_names}"
+ return True, reason
+
+
+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)
+
+ self._op = _sgl_fa3_op()
+ scheduler_packet = getattr(
+ torch.ops.sgl_kernel,
+ "get_scheduler_metadata",
+ None,
+ )
+ 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,
+ ]
+ args.append(0)
+ args.extend(
+ (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("sglang-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,
+ ]
+ args.append(0)
+ args.extend(
+ (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("sglang-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,
+ ]
+ args.append(0)
+ 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("sglang-kernel FA3 did not write to the supplied output")
+ return output
+
+
+__all__ = ["SglFa3DecodeKernel", "sgl_fa3_support"]
diff --git a/src/sparsevllm/kernels/external/sgl/moe.py b/src/sparsevllm/kernels/external/sgl/moe.py
new file mode 100644
index 00000000..dfb030d1
--- /dev/null
+++ b/src/sparsevllm/kernels/external/sgl/moe.py
@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+import importlib
+
+import torch
+import triton
+
+from sparsevllm.kernels.external.sgl.support import sgl_kernel_support
+from sparsevllm.kernels.moe import MoeAlignment
+
+
+def sgl_moe_alignment_support() -> tuple[bool, str]:
+ """Check the SGL expert-alignment API used by the Triton MoE provider."""
+
+ 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, (
+ "sglang-kernel MoE alignment failed to load: "
+ f"{type(error).__name__}: {error}"
+ )
+ return (True, reason) if callable(alignment) else (
+ False,
+ "sglang-kernel moe_align_block_size is not callable",
+ )
+
+
+def sgl_moe_align_block_size(
+ topk_ids: torch.Tensor,
+ *,
+ block_size: int,
+ num_experts: int,
+) -> MoeAlignment:
+ """Group local expert assignments with the SGL CUDA kernel."""
+
+ 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)
+ num_assignments = int(topk_ids.numel())
+ max_num_tokens_padded = triton.cdiv(
+ num_assignments + 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(
+ num_experts + 1,
+ dtype=torch.int32,
+ device=topk_ids.device,
+ )
+ from sgl_kernel import moe_align_block_size
+
+ # 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,
+ 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,
+ num_tokens_post_padded=num_tokens_post_padded,
+ block_size=int(block_size),
+ naive=False,
+ )
+
+
+__all__ = [
+ "sgl_moe_align_block_size",
+ "sgl_moe_alignment_support",
+]
diff --git a/src/sparsevllm/kernels/external/sgl/support.py b/src/sparsevllm/kernels/external/sgl/support.py
new file mode 100644
index 00000000..e26ca8d8
--- /dev/null
+++ b/src/sparsevllm/kernels/external/sgl/support.py
@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+import importlib
+import importlib.metadata
+import importlib.util
+import re
+
+_MIN_VERSION = (0, 4, 5)
+_MAX_VERSION = (0, 4, 6)
+_DISTRIBUTION = "sglang-kernel"
+
+
+def sgl_kernel_support(feature: str) -> tuple[bool, str]:
+ """Check the declared version and load the architecture-specific extension."""
+
+ try:
+ package_spec = importlib.util.find_spec("sgl_kernel")
+ except (ImportError, ValueError) as error:
+ return False, f"{_DISTRIBUTION} package discovery failed: {error}"
+ if package_spec is None:
+ return False, f"{_DISTRIBUTION} is not installed"
+ try:
+ version = importlib.metadata.version(_DISTRIBUTION)
+ except importlib.metadata.PackageNotFoundError:
+ 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 {_DISTRIBUTION}>=0.4.5,<0.4.6, got {version}"
+ try:
+ importlib.import_module("sgl_kernel")
+ except Exception as error:
+ return False, (
+ f"{_DISTRIBUTION} {version} {feature} failed to load: "
+ f"{type(error).__name__}: {error}"
+ )
+ return True, f"{_DISTRIBUTION} {version} {feature} is available"
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/kernels/tilelang/__init__.py b/src/sparsevllm/kernels/tilelang/__init__.py
new file mode 100644
index 00000000..4efc19f9
--- /dev/null
+++ b/src/sparsevllm/kernels/tilelang/__init__.py
@@ -0,0 +1 @@
+"""Repository-owned adapters for optional TileLang kernels."""
diff --git a/src/sparsevllm/kernels/tilelang/mla/LICENSE.tilelang b/src/sparsevllm/kernels/tilelang/mla/LICENSE.tilelang
new file mode 100644
index 00000000..eef5bb23
--- /dev/null
+++ b/src/sparsevllm/kernels/tilelang/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/kernels/tilelang/mla/README.md b/src/sparsevllm/kernels/tilelang/mla/README.md
new file mode 100644
index 00000000..7a067cd6
--- /dev/null
+++ b/src/sparsevllm/kernels/tilelang/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/kernels/tilelang/mla/__init__.py b/src/sparsevllm/kernels/tilelang/mla/__init__.py
new file mode 100644
index 00000000..8d349656
--- /dev/null
+++ b/src/sparsevllm/kernels/tilelang/mla/__init__.py
@@ -0,0 +1 @@
+"""GLM MLA TileLang kernels loaded lazily by the runtime adapter."""
diff --git a/src/sparsevllm/kernels/tilelang/mla/decode.py b/src/sparsevllm/kernels/tilelang/mla/decode.py
new file mode 100644
index 00000000..566c808b
--- /dev/null
+++ b/src/sparsevllm/kernels/tilelang/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/kernels/tilelang/mla/runtime.py b/src/sparsevllm/kernels/tilelang/mla/runtime.py
new file mode 100644
index 00000000..85152e80
--- /dev/null
+++ b/src/sparsevllm/kernels/tilelang/mla/runtime.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 collections.abc import Callable
+from dataclasses import dataclass
+from importlib import metadata
+
+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.kernels.tilelang.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
+
+ 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),)
+ ](
+ 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/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/kernels/triton/column_parallel_rmsnorm.py b/src/sparsevllm/kernels/triton/column_parallel_rmsnorm.py
new file mode 100644
index 00000000..9cfec1a8
--- /dev/null
+++ b/src/sparsevllm/kernels/triton/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/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 98%
rename from src/sparsevllm/triton_kernel/gqa_flash_decoding_stage1.py
rename to src/sparsevllm/kernels/triton/gqa_flash_decoding_stage1.py
index 140bc22d..15b3662e 100644
--- a/src/sparsevllm/triton_kernel/gqa_flash_decoding_stage1.py
+++ b/src/sparsevllm/kernels/triton/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/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 64%
rename from src/sparsevllm/triton_kernel/minimax_m2_router.py
rename to src/sparsevllm/kernels/triton/minimax_m2_router.py
index d0f7513f..3c613d74 100644
--- a/src/sparsevllm/triton_kernel/minimax_m2_router.py
+++ b/src/sparsevllm/kernels/triton/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.kernels.triton.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/kernels/triton/mla/LICENSE.lightllm b/src/sparsevllm/kernels/triton/mla/LICENSE.lightllm
new file mode 100644
index 00000000..261eeb9e
--- /dev/null
+++ b/src/sparsevllm/kernels/triton/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/kernels/triton/mla/README.md b/src/sparsevllm/kernels/triton/mla/README.md
new file mode 100644
index 00000000..c931206f
--- /dev/null
+++ b/src/sparsevllm/kernels/triton/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/kernels/triton/mla/__init__.py b/src/sparsevllm/kernels/triton/mla/__init__.py
new file mode 100644
index 00000000..0106939a
--- /dev/null
+++ b/src/sparsevllm/kernels/triton/mla/__init__.py
@@ -0,0 +1,44 @@
+"""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_MAX_WORKSPACE_CONFIG,
+ GLM_MLA_SOFTMAX_SCALE,
+ MlaDecodeLaunchConfig,
+ MlaDecodeWorkspace,
+ allocate_mla_decode_workspace,
+ 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
+from .decode_stage2 import decode_stage2
+from .gather_latent import (
+ gather_latent_history,
+ validate_gather_metadata,
+)
+
+__all__ = [
+ "DEFAULT_GLM_MLA_DECODE_CONFIG",
+ "GLM_MLA_MAX_WORKSPACE_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",
+ "select_glm_mla_decode_config",
+ "validate_copy_slot_mapping",
+ "validate_gather_metadata",
+ "validate_mla_decode_metadata",
+]
diff --git a/src/sparsevllm/kernels/triton/mla/copy_latent.py b/src/sparsevllm/kernels/triton/mla/copy_latent.py
new file mode 100644
index 00000000..286ed9c0
--- /dev/null
+++ b/src/sparsevllm/kernels/triton/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/kernels/triton/mla/decode_schedule.py b/src/sparsevllm/kernels/triton/mla/decode_schedule.py
new file mode 100644
index 00000000..db838f01
--- /dev/null
+++ b/src/sparsevllm/kernels/triton/mla/decode_schedule.py
@@ -0,0 +1,513 @@
+# 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
+
+import math
+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()
+
+# 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:
+ """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,
+ max_context_len: int | None = None,
+ valid_batch_size: int | None = None,
+) -> 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 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 "
+ f"{tuple(request_indices.shape)}"
+ )
+ 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()
+ real_request_rows: list[int] = []
+ for batch_index, (request_row, length) in enumerate(
+ zip(request_rows, lengths)
+ ):
+ if length < 0 or length > context_capacity:
+ raise ValueError(
+ f"context_lens[{batch_index}]={length} is outside "
+ f"[0, {context_capacity}]"
+ )
+ 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 batch_index < valid_batch_size:
+ 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,
+ *,
+ 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:
+ """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}"
+ )
+ 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,
+ 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,
+ q_rope,
+ latent_cache,
+ rope_cache,
+ active_slots,
+ request_indices,
+ context_lens,
+ 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,
+ 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/kernels/triton/mla/decode_stage1.py b/src/sparsevllm/kernels/triton/mla/decode_stage1.py
new file mode 100644
index 00000000..e4b824a9
--- /dev/null
+++ b/src/sparsevllm/kernels/triton/mla/decode_stage1.py
@@ -0,0 +1,499 @@
+# 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,
+ attn_score,
+ 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,
+ stride_score_batch,
+ stride_score_head,
+ stride_score_token,
+ 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,
+ 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)
+ 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,
+ )
+ raw_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,
+ )
+ 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,
+ -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,
+ *,
+ attn_score: torch.Tensor | None = None,
+ max_context_len: int | None = None,
+ 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,
+ }
+ 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:
+ 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 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):
+ 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
+ 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,
+ latent_cache,
+ rope_cache,
+ softmax_scale,
+ active_slots,
+ request_indices,
+ context_lens,
+ mid_output,
+ mid_logsumexp,
+ score_arg,
+ *active_slots.stride(),
+ *q_latent.stride(),
+ *q_rope.stride(),
+ *latent_cache.stride(),
+ *rope_cache.stride(),
+ *mid_output.stride(),
+ *mid_logsumexp.stride(),
+ *score_strides,
+ 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,
+ 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/kernels/triton/mla/decode_stage2.py b/src/sparsevllm/kernels/triton/mla/decode_stage2.py
new file mode 100644
index 00000000..b07aa5cf
--- /dev/null
+++ b/src/sparsevllm/kernels/triton/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/kernels/triton/mla/gather_latent.py b/src/sparsevllm/kernels/triton/mla/gather_latent.py
new file mode 100644
index 00000000..88bbcfa5
--- /dev/null
+++ b/src/sparsevllm/kernels/triton/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/src/sparsevllm/triton_kernel/moe.py b/src/sparsevllm/kernels/triton/moe.py
similarity index 89%
rename from src/sparsevllm/triton_kernel/moe.py
rename to src/sparsevllm/kernels/triton/moe.py
index cc0e2088..187d7ce0 100644
--- a/src/sparsevllm/triton_kernel/moe.py
+++ b/src/sparsevllm/kernels/triton/moe.py
@@ -1,30 +1,164 @@
from __future__ import annotations
-from dataclasses import dataclass
+from collections.abc import Callable
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,
+ 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,
+ 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(
@@ -1084,6 +1218,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 +1261,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/kernels/triton/moe_biased_sigmoid.py b/src/sparsevllm/kernels/triton/moe_biased_sigmoid.py
new file mode 100644
index 00000000..d141de3b
--- /dev/null
+++ b/src/sparsevllm/kernels/triton/moe_biased_sigmoid.py
@@ -0,0 +1,241 @@
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+
+_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,
+ 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 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,
+ *,
+ 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__ = ["fused_topk_biased_sigmoid", "topk_biased_sigmoid"]
diff --git a/src/sparsevllm/triton_kernel/moe_config.py b/src/sparsevllm/kernels/triton/moe_config.py
similarity index 86%
rename from src/sparsevllm/triton_kernel/moe_config.py
rename to src/sparsevllm/kernels/triton/moe_config.py
index d19837ca..e7de53d5 100644
--- a/src/sparsevllm/triton_kernel/moe_config.py
+++ b/src/sparsevllm/kernels/triton/moe_config.py
@@ -96,6 +96,77 @@ 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)
+_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(
+ 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 _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(
@@ -391,6 +462,17 @@ def _resolve_moe_gemm_config(
hidden_size=hidden_size,
intermediate_size=intermediate_size,
)
+ 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,
+ )
+ 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/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 73%
rename from src/sparsevllm/triton_kernel/omnikv_fused.py
rename to src/sparsevllm/kernels/triton/omnikv_fused.py
index 7fefd513..82499166 100644
--- a/src/sparsevllm/triton_kernel/omnikv_fused.py
+++ b/src/sparsevllm/kernels/triton/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/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 90%
rename from src/sparsevllm/triton_kernel/silu_and_mul.py
rename to src/sparsevllm/kernels/triton/silu_and_mul.py
index e3e7ef51..b174fec6 100644
--- a/src/sparsevllm/triton_kernel/silu_and_mul.py
+++ b/src/sparsevllm/kernels/triton/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/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/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/layers/attention.py b/src/sparsevllm/layers/attention.py
index 55731676..bce1d893 100644
--- a/src/sparsevllm/layers/attention.py
+++ b/src/sparsevllm/layers/attention.py
@@ -3,7 +3,10 @@
import torch
from torch import nn
+from sparsevllm.engine.cache_manager import ExplicitKVPayload
from sparsevllm.layers.attention_backend import TritonAttentionBackend
+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
@@ -56,6 +59,9 @@ def __init__(
head_dim,
scale,
num_kv_heads,
+ *,
+ prefill_op: PreparedPrefillAttentionOp | None = None,
+ decode_launch_op: PreparedDecodeAttentionLaunchOp | None = None,
):
super().__init__()
self.num_heads = num_heads
@@ -63,6 +69,8 @@ def __init__(
self.scale = scale
self.num_kv_heads = num_kv_heads
self.attention_backend = TritonAttentionBackend()
+ self.prefill_op = prefill_op
+ self.decode_launch_op = decode_launch_op
def forward(
self,
@@ -86,28 +94,57 @@ 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(
+ 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_op 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_op.run(
+ 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,
@@ -135,9 +172,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 +190,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: "
@@ -167,6 +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)
+ 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(
@@ -187,6 +244,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 6042cf74..02a93372 100644
--- a/src/sparsevllm/layers/attention_backend.py
+++ b/src/sparsevllm/layers/attention_backend.py
@@ -2,15 +2,19 @@
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
-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
@@ -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."""
@@ -100,6 +118,33 @@ class TritonAttentionBackend:
def __init__(self) -> None:
record_operator_binding("Attention", self)
+ 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,
@@ -109,30 +154,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)
+ 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,20 +191,20 @@ 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
- def _debug_check_prefill_bounds(
+ def debug_check_prefill_bounds(
self,
q: torch.Tensor,
view: PrefillComputeView,
@@ -168,39 +215,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 +261,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(
@@ -228,11 +277,15 @@ 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
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 +295,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,57 +336,59 @@ 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,
block_seq,
+ gqa_block_n,
+ gqa_num_warps,
)
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 +397,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 +410,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
+ from sparsevllm.kernels.triton.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 +487,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/src/sparsevllm/layers/embed_head.py b/src/sparsevllm/layers/embed_head.py
index 8df9a445..5128ca8b 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(y)
+ return self.parallel_context.tp_all_reduce(y) if self.reduce_results else y
class ParallelLMHead(VocabParallelEmbedding):
diff --git a/src/sparsevllm/layers/layernorm.py b/src/sparsevllm/layers/layernorm.py
index c4efbc25..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,
)
@@ -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.kernels.triton.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/layers/linear.py b/src/sparsevllm/layers/linear.py
index 018c9a03..995f0664 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(y)
+ return y
diff --git a/src/sparsevllm/layers/mla_attention.py b/src/sparsevllm/layers/mla_attention.py
new file mode 100644
index 00000000..d99fe078
--- /dev/null
+++ b/src/sparsevllm/layers/mla_attention.py
@@ -0,0 +1,1022 @@
+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,
+ MlaLatentWrite,
+ MlaLatentPayload,
+ PrefillComputeView,
+)
+from sparsevllm.layers.attention_backend import TritonAttentionBackend
+from sparsevllm.operators.mla_attention import (
+ MlaAttentionOpSpec,
+ MlaAttentionProvider,
+ resolve_mla_attention_provider,
+)
+from sparsevllm.kernels.triton.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_cu_seqlens: torch.Tensor
+ 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
+
+ @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
+ 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, ...]
+ total_visible_tokens: int
+ max_context_len: int
+ required_workspace_bytes: int
+
+ def matches(
+ self,
+ validation_scope: object,
+ meta: AttentionViewMeta,
+ cache_slot_count: int,
+ query_tokens: 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.source_query_tokens == int(query_tokens)
+ and self.cache_slot_count == int(cache_slot_count)
+ )
+
+
+@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
+ max_query_len: 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,
+ 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:
+ """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}.")
+ 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()
+ 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(
+ max(
+ kv_projection_phase_bytes,
+ attention_phase_bytes,
+ output_projection_phase_bytes,
+ )
+ + metadata_bytes
+ )
+
+
+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,
+ hidden_size: int,
+ projection_chunk_size: 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}."
+ )
+ 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, "
+ f"got {self.spec.qk_head_dim}/{self.spec.value_head_dim}."
+ )
+ 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(
+ cls,
+ *,
+ spec: MlaAttentionOpSpec,
+ 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,
+ device=device,
+ max_batch_size=max_batch_size,
+ )
+ return cls(
+ spec=spec,
+ provider=provider,
+ prefill_workspace_bytes=prefill_workspace_bytes,
+ hidden_size=hidden_size,
+ projection_chunk_size=projection_chunk_size,
+ )
+
+ @property
+ 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,
+ *,
+ 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,
+ query_tokens: 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,
+ query_tokens,
+ ):
+ 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 = _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}."
+ )
+ 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,
+ 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,
+ )
+ 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,
+ )
+ 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,
+ 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,
+ 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,
+ 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,
+ *,
+ query_tokens: int,
+ ) -> 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]),
+ query_tokens=query_tokens,
+ )
+ 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_cu_seqlens=plan.packed_cu_seqlens,
+ 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,
+ )
+
+ 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,
+ )
+
+ @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,
+ 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}."
+ )
+ 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),
+ ("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}."
+ )
+ 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,
+ 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,
+ 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,
+ context_lens=history.context_lens,
+ max_context_len=history.max_context_len,
+ ),
+ payload=ExplicitKVPayload(
+ k_cache=workset.expanded_k,
+ v_cache=workset.expanded_v,
+ metadata={
+ "layout": "mla_packed_varlen",
+ "cu_seqlens_k": history.packed_cu_seqlens,
+ },
+ ),
+ )
+
+ 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)
+ 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=context.attention_validation_scope,
+ 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",
+ "MlaPrefillWorkset",
+ "estimate_mla_prefill_workspace_bytes",
+]
diff --git a/src/sparsevllm/layers/packed_moe.py b/src/sparsevllm/layers/packed_moe.py
new file mode 100644
index 00000000..9a4d0f2e
--- /dev/null
+++ b/src/sparsevllm/layers/packed_moe.py
@@ -0,0 +1,308 @@
+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",
+ }
+ checkpoint_scale_dtype = torch.bfloat16
+
+ 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)
+ 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
+ ):
+ 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=scale_dtype,
+ ),
+ )
+ self.register_buffer(
+ "w2_scale_inv",
+ torch.empty(
+ self.num_local_experts,
+ self.hidden_size // 128,
+ self.intermediate_size // 128,
+ dtype=scale_dtype,
+ ),
+ )
+ 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 != self.checkpoint_scale_dtype:
+ raise TypeError(
+ 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:
+ 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/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..f2447f1e 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,7 +90,6 @@ class ModelRuntimeCompatibility:
decode_cuda_graph_methods=frozenset(CANONICAL_SPARSE_METHODS),
)
-
QWEN3_MOE_EP_COMPATIBILITY = ModelRuntimeCompatibility(
sparse_methods=_MOE_SPARSE_METHODS,
prefix_cache_methods=frozenset(
@@ -131,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=False,
+ decode_cuda_graph_methods=frozenset(
+ {"", "streamingllm", "snapkv", "h2o", "omnikv", "rkv"}
+ ),
+)
+
MODEL_RUNTIME_COMPATIBILITY = {
**{
(model_type, ParallelMode.STANDARD): DENSE_MODEL_COMPATIBILITY
@@ -142,6 +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_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
new file mode 100644
index 00000000..e041fd77
--- /dev/null
+++ b/src/sparsevllm/models/glm4_moe_lite.py
@@ -0,0 +1,1115 @@
+from __future__ import annotations
+
+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 (
+ ParallelContext,
+ get_parallel_context,
+)
+from sparsevllm.models.layout import resolve_attention_qk_head_dim
+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,
+ MergedReplicatedLinear,
+ 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.all_reduce import (
+ PreparedAllReduceOp,
+ 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,
+ resolve_moe_router_provider,
+)
+from sparsevllm.platforms import device_runtime
+from sparsevllm.utils.context import get_context
+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$"
+)
+_SHARED_EXPERT_SOURCE_RE = re.compile(
+ r"^model\.layers\.(\d+)\.mlp\.shared_experts\."
+ r"(gate_proj|up_proj|down_proj)\.weight$"
+)
+
+
+@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 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_parallel_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_parallel_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,
+ *,
+ device: torch.device | str,
+ 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."""
+
+ 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=resolve_attention_qk_head_dim(config),
+ 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(decode_cuda_graph),
+ )
+ return MLAAttention.bind(
+ spec=spec,
+ 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,
+ )
+
+
+class Glm4MoeLiteAttention(nn.Module):
+ """GLM projections around a shared latent-MLA execution object."""
+
+ def __init__(
+ self,
+ config: Glm4MoeLiteConfig,
+ 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)
+ 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 = 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:
+ 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.")
+
+ self.fused_qkv_a_proj = MergedReplicatedLinear(
+ int(config.hidden_size),
+ [
+ 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)
+ self.q_b_proj = ColumnParallelLinear(
+ self.q_lora_rank,
+ self.num_heads * self.qk_head_dim,
+ bias=False,
+ )
+ 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),
+ reduce_results=runtime_config is None,
+ )
+
+ 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 forward(
+ self,
+ positions: torch.Tensor,
+ hidden_states: torch.Tensor,
+ rotary_emb: RotaryEmbedding,
+ ) -> torch.Tensor:
+ 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,
+ )
+
+ 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.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:
+ 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):
+ 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)
+ )
+ 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.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)
+ 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
+
+
+class Glm4MoeLitePackedExperts(PackedMoeExperts):
+ def __init__(
+ self,
+ config: Glm4MoeLiteConfig,
+ *,
+ 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 = 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
+ )
+ packed_top_k = self.routed_top_k + int(self.fuses_shared_decode)
+ super().__init__(
+ num_experts=packed_num_experts,
+ hidden_size=int(config.hidden_size),
+ intermediate_size=int(config.moe_intermediate_size),
+ 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=parallel_context,
+ )
+ self.shared_expert_id = (
+ self.routed_num_experts if self.fuses_shared_decode else None
+ )
+ 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,
+ 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.")
+ 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__(
+ self,
+ config: Glm4MoeLiteConfig,
+ *,
+ 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(
+ f"mlp_chunk_size must be positive, got {self.mlp_chunk_size}."
+ )
+ self.gate = Glm4MoeLiteRouter(config)
+ self.experts = Glm4MoeLitePackedExperts(
+ config,
+ decode_cuda_graph=decode_cuda_graph,
+ )
+ 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=False,
+ activation_provider=resolve_silu_and_mul_provider(
+ activation_dtype=model_activation_dtype(config),
+ ),
+ )
+ )
+
+ 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(
+ "Glm4MoeLiteSparseMoeBlock expects [tokens, hidden], got "
+ f"{tuple(hidden_states.shape)}."
+ )
+ 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,
+ )
+ 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(local_output)
+ 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:
+ 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 device_runtime.is_stream_capturing()
+ else int(local_hit_count.item())
+ )
+ if debug_enabled:
+ # Preserve the general EP composition and routed-only debug
+ # 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.tp_size > 1:
+ shared = self.parallel_context.tp_all_reduce(shared)
+ output = routed + shared
+ elif self.parallel_context.ep_size > 1:
+ 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.
+ 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(local_output)
+ )
+ else:
+ # Retain the pure-EP semantic path for direct module use: the
+ # shared expert is replicated rather than TP-sharded.
+ 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)
+ 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(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 = 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(partials)
+ )
+ 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 synchronized shared-expert branch.
+ self.debug_last_output = output.detach().clone()
+ return output
+
+
+class Glm4MoeLiteDecoderLayer(nn.Module):
+ def __init__(
+ self,
+ config: Glm4MoeLiteConfig,
+ layer_idx: int,
+ mla_attention: MLAAttention,
+ *,
+ 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):
+ 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(mlp_chunk_size),
+ quantization=None,
+ reduce_results=runtime_config is None,
+ activation_provider=resolve_silu_and_mul_provider(
+ activation_dtype=model_activation_dtype(config),
+ ),
+ )
+ elif layer_type == "sparse":
+ self.mlp = Glm4MoeLiteSparseMoeBlock(
+ config,
+ mlp_chunk_size=mlp_chunk_size,
+ decode_cuda_graph=decode_cuda_graph,
+ runtime_config=runtime_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)
+ 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,
+ )
+ 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
+
+
+class Glm4MoeLiteModel(nn.Module):
+ def __init__(
+ self,
+ config: Glm4MoeLiteConfig,
+ mla_attention: MLAAttention,
+ *,
+ 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(
+ 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,
+ 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))
+ ]
+ )
+ 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:
+ 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
+ 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
+
+
+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),
+ }
+
+ @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,
+ }
+ 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,
+ *,
+ mla_attention: MLAAttention,
+ mlp_chunk_size: int,
+ decode_cuda_graph: 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:
+ self.lm_head.weight.data = self.model.embed_tokens.weight.data
+ 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:
+ self.runtime_config.close()
+
+ 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
+
+ def map_weight_name(self, source_weight_name: str) -> str | 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
+ layer_idx, global_expert_id, projection = match.groups()
+ 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}."
+ 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()
+ return WeightTarget(
+ self._sparse_block(int(layer_idx)).experts,
+ (int(global_expert_id), projection),
+ )
+
+ 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
+ expert_id, projection = target.shard_id
+ target.module.load_expert_weight(
+ expert_id,
+ projection,
+ loaded_weight,
+ loaded_scale,
+ )
+ return 1
+
+ def validate_loaded_weights(self, loaded_parameter_names: set[str]) -> None:
+ packed_experts = {
+ name
+ for name, _ in self.named_parameters()
+ if name.endswith((".mlp.experts.w13_weight", ".mlp.experts.w2_weight"))
+ }
+ 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()
+
+ @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 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",
+ "build_glm4_moe_lite_mla_attention",
+]
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/minimax_m2.py b/src/sparsevllm/models/minimax_m2.py
index 65cd7404..1153a16f 100644
--- a/src/sparsevllm/models/minimax_m2.py
+++ b/src/sparsevllm/models/minimax_m2.py
@@ -1,29 +1,43 @@
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,
+ get_parallel_context,
+)
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 (
+ PreparedAllReduceOp,
+ prepare_parallel_all_reduce,
+)
+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
from sparsevllm.utils.log import logger
from sparsevllm.utils.weight_target import WeightTarget
@@ -39,6 +53,104 @@
)
+@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 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_parallel_all_reduce(
+ parallel_context.world,
+ max_rows=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_parallel_all_reduce(
+ parallel_context.attention,
+ max_rows=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__()
@@ -51,7 +163,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,
)
@@ -70,198 +182,61 @@ 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, config) -> 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:
+ def __init__(
+ self,
+ config,
+ *,
+ cuda_graph: bool | None = None,
+ ) -> None:
+ 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,
- cuda_graph=bool(getattr(config, "decode_cuda_graph", False)),
- tp_size=self.tp_size,
+ fp8_enabled=True,
+ cuda_graph=(
+ bool(getattr(config, "decode_cuda_graph", False))
+ if cuda_graph is None
+ else bool(cuda_graph)
+ ),
routing_method="biased_sigmoid",
scale_dtype=torch.float32,
+ model_label="MiniMax",
+ provider_resolver=resolve_moe_provider,
+ parallel_context=get_parallel_context(),
)
- 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(
+class MiniMaxM2SparseMoeBlock(nn.Module):
+ def __init__(
self,
- global_expert_id: int,
- projection: str,
- loaded_weight: torch.Tensor,
- loaded_scale: torch.Tensor | None,
+ config,
+ runtime_config: MiniMaxM2RuntimeConfig | 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 {"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,
- )
-
-
-class MiniMaxM2SparseMoeBlock(nn.Module):
- def __init__(self, config) -> 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:
@@ -287,13 +262,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)
@@ -318,6 +302,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,
@@ -341,6 +326,12 @@ def __init__(self, config) -> None:
self.head_dim,
self.head_dim**-0.5,
self.num_kv_heads,
+ 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(
@@ -365,15 +356,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,
@@ -405,8 +405,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,
@@ -421,15 +429,44 @@ class MiniMaxM2ForCausalLM(nn.Module):
"v_proj": ("qkv_proj", "v"),
}
- def __init__(self, config) -> None:
+ @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,
+ 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)
@@ -640,12 +677,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={} attention TP {}/{} MoE TP "
+ "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,
+ 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/models/qwen3.py b/src/sparsevllm/models/qwen3.py
index 9f467fe1..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
@@ -169,6 +170,8 @@ def __init__(
hidden_act: str,
mlp_chunk_size: int = 16384,
quantization=None,
+ reduce_results: bool = True,
+ activation_provider: SiluAndMulProvider | None = None,
) -> None:
super().__init__()
self.gate_up_proj = MergedColumnParallelLinear(
@@ -182,9 +185,10 @@ def __init__(
hidden_size,
bias=False,
quantization=quantization,
+ 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/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 3a9d04d5..a438169c 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
@@ -43,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))
@@ -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/src/sparsevllm/models/spec.py b/src/sparsevllm/models/spec.py
index 439ef192..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, ...] = ()
@@ -159,6 +160,17 @@ 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",
+ 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"),
+ top_k_field="num_experts_per_tok",
+ ),
}
)
diff --git a/src/sparsevllm/operators/activation.py b/src/sparsevllm/operators/activation.py
new file mode 100644
index 00000000..072c9c00
--- /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.kernels.triton.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/all_reduce.py b/src/sparsevllm/operators/all_reduce.py
index ee848e3d..7acfa046 100644
--- a/src/sparsevllm/operators/all_reduce.py
+++ b/src/sparsevllm/operators/all_reduce.py
@@ -1,84 +1,572 @@
from __future__ import annotations
-from typing import Protocol
+import ctypes
+import re
+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
-from sparsevllm.utils.log import logger
+import sparsevllm.platforms as platforms
+from sparsevllm.operators.registry import (
+ OpRegistry,
+ OpResolver,
+ SupportResult,
+ runtime_version_at_least,
+)
+from sparsevllm.platforms.interface import DeviceCaps, PlatformEnum
+if TYPE_CHECKING:
+ from sparsevllm.distributed.parallel_context import ParallelGroup
-class AllReduceProvider(Protocol):
- name: str
- def run(self, tensor: torch.Tensor) -> torch.Tensor: ...
+@dataclass(frozen=True)
+class AllReduceOpSpec:
+ world_size: 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_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 TorchDistributedAllReduceProvider:
- name = "torch_distributed"
- def __init__(self, group: dist.ProcessGroup | None) -> None:
- self.group = group
+class AllReduceProvider:
+ name = ""
+ priority = 0
- def run(self, tensor: torch.Tensor) -> torch.Tensor:
- dist.all_reduce(tensor, group=self.group)
- return tensor
+ def prepare(
+ self,
+ spec: AllReduceOpSpec,
+ *,
+ group: dist.ProcessGroup | None,
+ rank: int,
+ device_index: int | None = None,
+ ) -> None:
+ del spec, group, rank, device_index
+ def close(self) -> None:
+ pass
-class HopperTp2FlashInferAllReduceProvider:
- name = "hopper_tp2_flashinfer"
- hidden_size = 2048
- max_rows = 256
+ def run(
+ self,
+ spec: AllReduceOpSpec,
+ tensor: torch.Tensor,
+ *,
+ group: dist.ProcessGroup | None,
+ ) -> torch.Tensor:
+ """Return the reduced tensor, which may or may not alias the input."""
+ raise NotImplementedError
+
+
+ALL_REDUCE_REGISTRY: OpRegistry[AllReduceOpSpec, AllReduceProvider] = OpRegistry(
+ "all-reduce"
+)
+
+
+@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"
+ 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.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 a profiled BF16 topology/shape, got "
+ f"world_size={spec.world_size} hidden_size={spec.hidden_size} "
+ f"dtype={spec.dtype}"
+ )
+ if spec.max_rows > profile.max_rows:
+ return SupportResult.no(
+ f"requires max_rows <= {profile.max_rows}, got {spec.max_rows}"
+ )
+ 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 __init__(self, group: dist.ProcessGroup) -> None:
- from flashinfer import comm
+ def prepare(self, spec, *, group, rank, device_index=None) -> None:
+ from flashinfer.comm import create_allreduce_fusion_workspace
- self.comm = comm
- self.fallback = TorchDistributedAllReduceProvider(group)
- self.workspace = comm.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.")
+ 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
+ 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}."
+ )
+ 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=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=profile.max_rows,
+ hidden_dim=spec.hidden_size,
+ dtype=spec.dtype,
group=group,
)
+ try:
+ 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 _supports(self, tensor: torch.Tensor) -> bool:
+ 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._profile is None:
+ raise RuntimeError("FlashInfer all-reduce provider was not prepared.")
+ if not tensor.is_cuda:
+ raise ValueError(
+ f"FlashInfer all-reduce requires a CUDA tensor, got {tensor.device}."
+ )
+ from flashinfer.comm import allreduce_fusion
+ from flashinfer.comm.trtllm_ar import AllReduceFusionPattern
+
+ 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 (
- tensor.is_cuda
- and tensor.dtype == torch.bfloat16
- and tensor.is_contiguous()
- and tensor.ndim >= 2
- and tensor.shape[-1] == self.hidden_size
- and tensor.numel() <= self.max_rows * self.hidden_size
+ SupportResult.no(
+ "FlashInfer communication APIs are missing: " + ", ".join(missing)
+ )
+ if missing
+ else SupportResult.yes()
)
- 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),
- workspace=self.workspace,
- pattern=self.comm.AllReduceFusionPattern.kAllReduce,
- launch_with_pdl=True,
- trigger_completion_at_end=tensor.numel() > 16 * self.hidden_size,
+ 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.view_as(tensor)
+ 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):
+ 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:
+ 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 _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 [..., hidden], got shape={tuple(tensor.shape)}."
+ )
+ 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 row count is outside the prepared range: "
+ f"rows={row_count} max_rows={spec.max_rows}."
+ )
+ 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(
- group: dist.ProcessGroup | None,
- world_size: int,
+ spec: AllReduceOpSpec,
+ *,
+ device_index: int | None = None,
) -> 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)
+ 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(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)
+
+
+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/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/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
new file mode 100644
index 00000000..5ce6fe47
--- /dev/null
+++ b/src/sparsevllm/operators/mla_attention.py
@@ -0,0 +1,805 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+
+import sparsevllm.platforms as platforms
+from sparsevllm.engine.cache_manager.base import (
+ DecodeComputeView,
+ ExplicitKVPayload,
+ MlaLatentPayload,
+ PrefillComputeView,
+)
+from sparsevllm.kernels.external.sgl.fa3 import SglFa3DecodeKernel, sgl_fa3_support
+from sparsevllm.kernels.tilelang.mla.runtime import (
+ TileMlaDecodeKernel,
+ tilelang_mla_support,
+)
+from sparsevllm.kernels.triton.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,
+)
+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
+_GLM_MLA_ROPE_DIM = 64
+_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)
+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
+ supports_explicit_prefill = False
+
+ def run(
+ self,
+ q_nope_absorbed: torch.Tensor,
+ q_rope: torch.Tensor,
+ view: DecodeComputeView,
+ output: torch.Tensor,
+ *,
+ validation_scope: object | None = None,
+ valid_batch_size: int | None = None,
+ ) -> torch.Tensor:
+ raise NotImplementedError
+
+
+MLA_ATTENTION_REGISTRY: OpRegistry[
+ MlaAttentionOpSpec,
+ MlaAttentionProvider,
+] = OpRegistry("MLA attention")
+
+
+@MLA_ATTENTION_REGISTRY.register
+class MlaTritonProvider(MlaAttentionProvider):
+ """Validated SM90 provider with caller-independent decode workspace."""
+
+ name = "triton_sm90"
+ priority = 100
+
+ def __init__(
+ self,
+ *,
+ op_spec: MlaAttentionOpSpec,
+ device: torch.device | str,
+ max_batch_size: int,
+ launch_config: MlaDecodeLaunchConfig | None = None,
+ ) -> 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._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=workspace_config,
+ )
+ self.device = self.workspace.block_size.device
+ self._validated_decode_metadata: tuple[
+ object,
+ list[
+ tuple[
+ torch.Tensor,
+ torch.Tensor,
+ torch.Tensor,
+ int,
+ int | None,
+ ]
+ ],
+ ] | None = None
+
+ @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 not in _VALIDATED_SM90_NAMES:
+ return SupportResult.no(
+ "requires validated H100 80GB HBM3 or H20 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.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}"
+ )
+ 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 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(
+ 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
+
+ def _validate_metadata(
+ self,
+ view: DecodeComputeView | PrefillComputeView,
+ payload: MlaLatentPayload,
+ *,
+ validation_scope: object | None,
+ valid_batch_size: int | None,
+ ) -> None:
+ cache_slot_count = int(payload.latent_cache.shape[0])
+ metadata_key = (
+ view.meta.active_slots,
+ view.meta.req_indices,
+ view.meta.context_lens,
+ cache_slot_count,
+ view.meta.max_context_len,
+ valid_batch_size,
+ )
+ 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 []
+ )
+ 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]
+ and entry[5] == metadata_key[5]
+ for entry in 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,
+ valid_batch_size=valid_batch_size,
+ )
+ 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,
+ valid_batch_size: int | 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,
+ valid_batch_size=valid_batch_size,
+ )
+ 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,
+ 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,
+ attn_score=view.meta.attn_score,
+ max_context_len=view.meta.max_context_len,
+ 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_sm90"
+ 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,
+ valid_batch_size: int | 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,
+ valid_batch_size=valid_batch_size,
+ )
+ payload = self._validate_run_inputs(
+ q_nope_absorbed,
+ q_rope,
+ view,
+ output,
+ )
+ self._validate_metadata(
+ view,
+ payload,
+ validation_scope=validation_scope,
+ valid_batch_size=valid_batch_size,
+ )
+ 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,
+ # Zero enables FA3's measured context-aware split heuristic.
+ num_splits=0,
+ 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}."
+ )
+ 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,
+ 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,
+ )
+
+
+@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
+ 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)
+
+ @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,
+ valid_batch_size: int | 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,
+ 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
+ # 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,
+ valid_batch_size=valid_batch_size,
+ )
+ payload = self._validate_run_inputs(
+ q_nope_absorbed,
+ q_rope,
+ view,
+ output,
+ )
+ self._validate_metadata(
+ view,
+ payload,
+ validation_scope=validation_scope,
+ valid_batch_size=valid_batch_size,
+ )
+ 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,
+ *,
+ device: torch.device | str,
+ max_batch_size: int,
+ launch_config: MlaDecodeLaunchConfig | None = None,
+) -> 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",
+ "MlaSglFa3Provider",
+ "MlaTileLangScoreProvider",
+ "MlaTritonProvider",
+ "resolve_mla_attention_provider",
+]
diff --git a/src/sparsevllm/operators/moe.py b/src/sparsevllm/operators/moe.py
index 5ab1ae7e..74eec8ba 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,
@@ -172,14 +173,165 @@ 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,
+ *,
+ 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,
+ )
+
+
+@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 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"),
+ (64, 32, 2048, 1536, 4, 1, 2, "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.kernels.external.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.kernels.triton.moe import fused_moe
+
+ alignment_impl = None
+ num_tokens = int(hidden_states.shape[0])
+ if (
+ num_tokens <= 64
+ 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(
+ hidden_states,
+ w13_weight,
+ w2_weight,
+ topk_ids,
+ topk_weights,
+ num_experts=spec.num_experts,
+ local_expert_start=local_expert_start,
+ alignment_impl=alignment_impl,
+ )
+
@MOE_REGISTRY.register
class TritonMinimaxM2FusedMoeProvider(MoeProvider):
name = "triton_minimax_m2_fused"
- priority = 110
+ priority = 125
gate_up_order = "gate_up"
@classmethod
@@ -261,7 +413,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,
)
@@ -345,7 +497,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
@@ -420,7 +572,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,
@@ -495,7 +647,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,
@@ -509,7 +661,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,
@@ -527,7 +679,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(
{
@@ -604,7 +756,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,
@@ -623,7 +775,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/moe_router.py b/src/sparsevllm/operators/moe_router.py
index fc8e74c3..19e4d457 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,8 +92,13 @@ 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]:
- from sparsevllm.triton_kernel.moe_topk import topk_softmax
+ 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.kernels.triton.moe_topk import topk_softmax
return topk_softmax(
router_logits,
@@ -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.kernels.triton.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/operators/prefill_attention.py b/src/sparsevllm/operators/prefill_attention.py
new file mode 100644
index 00000000..76cea45a
--- /dev/null
+++ b/src/sparsevllm/operators/prefill_attention.py
@@ -0,0 +1,440 @@
+from __future__ import annotations
+
+import math
+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,
+ *,
+ device_index: int | None = None,
+ ) -> None:
+ del spec, device_index
+
+ def close(self) -> None:
+ pass
+
+ 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
+
+
+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")
+
+
+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,
+ *,
+ 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,
+ spec,
+ q,
+ view,
+ *,
+ qo_indptr,
+ chunk_lens,
+ max_context_len,
+ layer_idx,
+ ):
+ del chunk_lens
+ payload, meta = _view_parts(view)
+ _validate_token_page_table(meta)
+ if self._state is None:
+ self.prepare(spec, device_index=q.device.index)
+ 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
+
+
+@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}")
+ 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(
+ self,
+ spec,
+ q,
+ view,
+ *,
+ qo_indptr,
+ chunk_lens,
+ max_context_len,
+ layer_idx,
+ ):
+ del spec, layer_idx
+ payload, meta = _view_parts(view)
+ _validate_token_page_table(meta)
+ from sparsevllm.kernels.triton.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
+
+
+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/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/src/sparsevllm/utils/context.py b/src/sparsevllm/utils/context.py
index 5ea917a9..a53d850c 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
@@ -29,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/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/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_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_attention_cache_storage.py b/tests/test_attention_cache_storage.py
new file mode 100644
index 00000000..9333c967
--- /dev/null
+++ b/tests/test_attention_cache_storage.py
@@ -0,0 +1,564 @@
+from __future__ import annotations
+
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import numpy as np
+import pytest
+import torch
+
+from sparsevllm.config import 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.snapkv import SnapKVCacheManager
+from sparsevllm.engine.cache_manager.storage import (
+ CacheLayout,
+ ExplicitKVStorage,
+ MlaLatentStorage,
+ create_attention_cache_storage,
+)
+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,
+ )
+
+@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),
+ ),
+ )
+
+
+@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,
+ 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_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,
+ 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_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)],
+)
+def test_standard_manager_accounts_storage_tensors_explicitly(
+ attention_tp_size,
+ local_heads,
+):
+ 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,
+ 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
+ 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["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",
+ "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])
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_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_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_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_dependency_constraints.py b/tests/test_dependency_constraints.py
index d3df2451..1e3187e0 100644
--- a/tests/test_dependency_constraints.py
+++ b/tests/test_dependency_constraints.py
@@ -1,18 +1,118 @@
+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 "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 "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
+ )
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_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 config["project"]["optional-dependencies"] == {
+ "cu129": [
+ "torch==2.11.0",
+ "flashinfer-python[cu12]>=0.6.15,<0.7",
+ "flashinfer-jit-cache>=0.6.15,<0.7",
+ ],
+ "cu130": [
+ "torch==2.11.0",
+ "flashinfer-python[cu13]>=0.6.15,<0.7",
+ "flashinfer-jit-cache>=0.6.15,<0.7",
+ ],
+ }
+ assert uv_config["sources"] == {
+ "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",
+ "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"]
+ dependencies = set(project["dependencies"])
+
+ 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_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_glm4_moe_lite.py b/tests/test_glm4_moe_lite.py
new file mode 100644
index 00000000..c2b35284
--- /dev/null
+++ b/tests/test_glm4_moe_lite.py
@@ -0,0 +1,715 @@
+from __future__ import annotations
+
+import os
+from contextlib import ExitStack, contextmanager
+from types import SimpleNamespace
+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,
+)
+
+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.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,
+)
+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
+
+
+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
+ 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 _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,
+ 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"),
+ hidden_size=64,
+ projection_chunk_size=8,
+ )
+
+
+@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(
+ 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)
+ 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),
+ mlp_chunk_size=config.mlp_chunk_size,
+ decode_cuda_graph=config.decode_cuda_graph,
+ )
+
+
+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_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,
+ "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)
+ 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),
+ projection_chunk_size=config.mlp_chunk_size,
+ )
+
+ 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_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()
+ with _construction_context(context):
+ 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)
+ 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)
+ 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
+
+ def torch_topk(_spec, logits, correction_bias, *, routed_scaling_factor):
+ scores = logits.sigmoid()
+ ids = torch.topk(scores + correction_bias, 4, dim=-1, sorted=False).indices
+ weights = scores.gather(1, ids)
+ weights = weights / (weights.sum(dim=-1, keepdim=True) + 1e-20)
+ return weights * routed_scaling_factor, ids
+
+ router.provider.run = 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(),
+ mlp_chunk_size=config.mlp_chunk_size,
+ decode_cuda_graph=config.decode_cuda_graph,
+ )
+ 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
+
+
+@pytest.mark.parametrize(("ep_rank", "ep_size"), [(1, 2), (3, 4)])
+def test_glm_ep_loader_keeps_only_local_experts(
+ 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,
+ )
+ 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
+
+
+@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.runtime_config = None
+ 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
+
+
+@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),
+ 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.runtime_config = None
+ 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,
+ 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)
+ 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:
+ 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.runtime_config = None
+ 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_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)
+ nn.Module.__init__(block)
+ block.parallel_context = context
+ block.runtime_config = None
+ 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"}),
+ 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)
+
+ 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
+ layer.runtime_config = None
+
+ 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
diff --git a/tests/test_glm_cuda_graph.py b/tests/test_glm_cuda_graph.py
new file mode 100644
index 00000000..b8d0df5e
--- /dev/null
+++ b/tests/test_glm_cuda_graph.py
@@ -0,0 +1,1359 @@
+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.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
+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="",
+ 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,
+ )
+ 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="",
+ 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,
+ 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..51efadc4
--- /dev/null
+++ b/tests/test_glm_runtime_compatibility.py
@@ -0,0 +1,276 @@
+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"),
+ ],
+)
+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
diff --git a/tests/test_h2o_cache_manager.py b/tests/test_h2o_cache_manager.py
index bddf13e2..a182c650 100644
--- a/tests/test_h2o_cache_manager.py
+++ b/tests/test_h2o_cache_manager.py
@@ -9,12 +9,15 @@
import torch
from sparsevllm.engine.cache_manager.base import (
+ AttentionViewMeta,
CacheManager,
+ ExplicitKVPayload,
LayerBatchStates,
PrefillComputeView,
)
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
@@ -360,12 +363,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 +409,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])
@@ -792,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_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_model.py b/tests/test_minimax_m2_model.py
index 28fb8599..058daf75 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,66 @@ 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 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):
@@ -427,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_minimax_m2_router.py b/tests/test_minimax_m2_router.py
index 8caba566..c3d281a9 100644
--- a/tests/test_minimax_m2_router.py
+++ b/tests/test_minimax_m2_router.py
@@ -4,10 +4,13 @@
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.kernels.triton.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")
diff --git a/tests/test_mla_attention_layer.py b/tests/test_mla_attention_layer.py
new file mode 100644
index 00000000..0d67137e
--- /dev/null
+++ b/tests/test_mla_attention_layer.py
@@ -0,0 +1,668 @@
+from __future__ import annotations
+
+from types import SimpleNamespace
+from unittest.mock import Mock, patch
+
+import pytest
+import torch
+
+import sparsevllm.layers.mla_attention as mla_attention_module
+from sparsevllm.engine.cache_manager import (
+ AttentionKeyComputeView,
+ AttentionViewMeta,
+ DecodeComputeView,
+ 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 get_context, reset_context, set_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)
+
+
+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,
+ 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,
+ hidden_size=64,
+ projection_chunk_size=8,
+ )
+
+
+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 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
+ 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,
+ 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
+ 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 == 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:
+ 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, query_tokens=1)
+ 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, query_tokens=2)
+ 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,
+ hidden_size=64,
+ projection_chunk_size=8,
+ )
+
+ assert attention.provider is provider
+ resolve.assert_called_once_with(
+ spec,
+ device="cpu",
+ max_batch_size=8,
+ )
+
+
+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_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)
+ 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, query_tokens=2)
+ second = attention.prepare_prefill_history(second_view, query_tokens=2)
+ reset_context()
+ attention.prepare_prefill_history(second_view, query_tokens=2)
+
+ 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
+
+
+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, query_tokens=2)
+ 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
+
+
+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.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
+ 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)
+ 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",
+)
+
+
+@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, 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")
+ 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"),
+ ),
+ query_tokens=6,
+ )
+ 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"),
+ ),
+ query_tokens=4,
+ )
+ 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
new file mode 100644
index 00000000..4672e38e
--- /dev/null
+++ b/tests/test_mla_attention_operator.py
@@ -0,0 +1,578 @@
+from __future__ import annotations
+
+from unittest.mock import Mock, patch
+
+import pytest
+import torch
+
+from sparsevllm.engine.cache_manager import (
+ AttentionViewMeta,
+ DecodeComputeView,
+ ExplicitKVPayload,
+ MlaLatentPayload,
+ PrefillComputeView,
+)
+from sparsevllm.operators.mla_attention import (
+ MLA_ATTENTION_REGISTRY,
+ MlaAttentionOpSpec,
+ MlaSglFa3Provider,
+ MlaTileLangScoreProvider,
+ MlaTritonProvider,
+)
+from sparsevllm.operators.registry import OpResolver
+from sparsevllm.platforms import DeviceCaps, PlatformEnum
+from sparsevllm.kernels.triton.mla import (
+ GLM_MLA_MAX_WORKSPACE_CONFIG,
+ 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)
+
+
+@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 (
+ patch(
+ "sparsevllm.operators.mla_attention.sgl_fa3_support",
+ return_value=(False, "sglang-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(device_name=device_name),
+ op_spec=spec,
+ device="cpu",
+ max_batch_size=8,
+ )
+
+ assert isinstance(resolved.provider, MlaTritonProvider)
+ assert resolved.rejected == (
+ ("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=spec.local_q_heads,
+ device=torch.device("cpu"),
+ config=GLM_MLA_MAX_WORKSPACE_CONFIG,
+ )
+
+
+@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 (
+ patch(
+ "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,
+ ),
+ patch("sparsevllm.operators.mla_attention.SglFa3DecodeKernel"),
+ ):
+ resolved = OpResolver(MLA_ATTENTION_REGISTRY).resolve(
+ spec,
+ _h100_caps(device_name=device_name),
+ op_spec=spec,
+ device="cpu",
+ max_batch_size=8,
+ )
+
+ assert type(resolved.provider) is MlaSglFa3Provider
+ assert resolved.rejected == (
+ (
+ "tilelang_score_sgl_fa3_h100",
+ tilelang_reason,
+ ),
+ )
+
+
+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.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,
+ _h100_caps(),
+ op_spec=spec,
+ device="cpu",
+ max_batch_size=8,
+ )
+
+ assert type(resolved.provider) is MlaTileLangScoreProvider
+ assert resolved.rejected == ()
+
+
+@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"),
+ (
+ {"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"),
+ (
+ {"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"),
+ ],
+)
+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_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)
+ 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)
+ validation_scope = object()
+
+ 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,
+ patch(
+ "sparsevllm.operators.mla_attention.validate_mla_decode_metadata"
+ ) as validate,
+ ):
+ 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()
+ 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,
+ max_context_len=None,
+ valid_batch_size=None,
+ )
+ assert kernel.call_count == 3
+ kernel.assert_called_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,
+ 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)
+ 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
new file mode 100644
index 00000000..5cfd01f8
--- /dev/null
+++ b/tests/test_mla_kernels.py
@@ -0,0 +1,854 @@
+from __future__ import annotations
+
+import ast
+from dataclasses import replace
+from pathlib import Path
+
+import pytest
+import torch
+
+from sparsevllm.kernels.triton.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,
+ select_glm_mla_decode_config,
+ validate_mla_decode_metadata,
+)
+
+
+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 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,
+ 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"
+ / "kernels"
+ / "triton"
+ / "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,
+ softmax_scale=GLM_MLA_SOFTMAX_SCALE,
+ )
+ 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}"
+ ),
+ )
+
+
+@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)
+ 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,
+ softmax_scale=GLM_MLA_SOFTMAX_SCALE,
+ )
+ 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,
+ softmax_scale=GLM_MLA_SOFTMAX_SCALE,
+ )
+
+
+@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)
+ 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,
+ softmax_scale=GLM_MLA_SOFTMAX_SCALE,
+ )
+ 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,
+ softmax_scale=GLM_MLA_SOFTMAX_SCALE,
+ )
diff --git a/tests/test_moe_config.py b/tests/test_moe_config.py
index a60cf648..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,
@@ -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,73 @@ 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
+
+
+@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_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_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_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_openai_api_server.py b/tests/test_openai_api_server.py
index 5bb2cce9..72b65563 100644
--- a/tests/test_openai_api_server.py
+++ b/tests/test_openai_api_server.py
@@ -29,7 +29,7 @@ def decode(self, token_ids, skip_special_tokens=True):
class _TransformersResponseTokenizer:
response_template = None
- def __init__(self, *, xml_tools=False, minimax_tools=False):
+ def __init__(self, *, xml_tools=False, minimax_tools=False, glm_tools=False):
self.chat_template = ""
if xml_tools:
self.chat_template += ""
@@ -38,6 +38,10 @@ def __init__(self, *, xml_tools=False, minimax_tools=False):
''
''
)
+ 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
@@ -76,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
@@ -115,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(
@@ -138,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
@@ -198,6 +218,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
@@ -2593,6 +2619,7 @@ def exit(self):
completion_top_logprobs=[None],
detokenizer=detokenizer,
emitted_text_len=1,
+ emitted_raw_text_len=1,
)
}
try:
@@ -2605,6 +2632,150 @@ 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),
+ 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["finish_reason"], "stop")
+ 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")
+ 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
from sparsevllm.entrypoints.openai.detokenizer import IncrementalDetokenizer
@@ -3110,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:]))
@@ -3121,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
@@ -3531,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")])
@@ -3626,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
@@ -3660,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)
@@ -5470,6 +5698,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
diff --git a/tests/test_operator_providers.py b/tests/test_operator_providers.py
index f6357d45..29fced49 100644
--- a/tests/test_operator_providers.py
+++ b/tests/test_operator_providers.py
@@ -5,7 +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,
@@ -23,7 +29,9 @@
FlashInferCutlassFp8MoeProvider,
HopperQwen36HybridFp8MoeProvider,
MoeOpSpec,
+ SglAlignedTritonGlmMoeProvider,
resolve_moe_provider,
+ use_packed_shared_experts,
)
from sparsevllm.operators.registry import OpResolver
from sparsevllm.platforms import DeviceCaps, PlatformEnum
@@ -123,16 +131,106 @@ def _gate_up_spec(**overrides) -> GateUpSwiGLUOpSpec:
return GateUpSwiGLUOpSpec(**values)
-def test_flashinfer_all_reduce_falls_back_before_unsupported_shape_launch():
- provider = HopperTp2FlashInferAllReduceProvider.__new__(
- HopperTp2FlashInferAllReduceProvider
+def _all_reduce_spec(**overrides) -> AllReduceOpSpec:
+ values = {
+ "world_size": 4,
+ "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)
+
+
+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 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(),
)
- provider.fallback = Mock()
- tensor = torch.randn(1, 248320, dtype=torch.bfloat16)
- provider.fallback.run.return_value = tensor
+ 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 provider.run(tensor) is tensor
- provider.fallback.run.assert_called_once_with(tensor)
+ 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])
@@ -192,6 +290,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",
[
@@ -453,6 +567,172 @@ def test_hopper_fused_moe_uses_profiled_tp_ep_shape():
assert resolved.provider.name == "triton_hopper_fused"
+@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,
+ 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.kernels.external.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=device_name,
+ ),
+ )
+
+ assert resolved.provider.name == "sgl_aligned_triton_glm"
+
+
+@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,
+ 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.kernels.external.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=device_name,
+ ),
+ )
+
+ 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
+
+
+@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,
+ 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.kernels.external.sgl.moe.sgl_moe_alignment_support",
+ return_value=(True, "supported"),
+ ):
+ resolved = OpResolver(MOE_REGISTRY).resolve(
+ spec,
+ _cuda_caps(
+ (9, 0),
+ native_fp8=False,
+ device_name=device_name,
+ ),
+ )
+
+ assert resolved.provider.name == "sgl_aligned_triton_glm"
+
+
+@pytest.mark.parametrize(
+ ("num_tokens", "expects_sgl_alignment"),
+ [(1, True), (2, True), (4, True), (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.kernels.triton.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)],
@@ -793,7 +1073,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
)
},
@@ -854,7 +1134,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_packed_moe.py b/tests/test_packed_moe.py
new file mode 100644
index 00000000..e5d3bb11
--- /dev/null
+++ b/tests/test_packed_moe.py
@@ -0,0 +1,71 @@
+from types import SimpleNamespace
+
+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
+
+
+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_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:
+ 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
diff --git a/tests/test_parallel_context.py b/tests/test_parallel_context.py
index e46c67a3..eab58a6a 100644
--- a/tests/test_parallel_context.py
+++ b/tests/test_parallel_context.py
@@ -12,13 +12,15 @@
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.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
@@ -231,6 +233,28 @@ def test_ep_broadcast_rejects_invalid_source_rank():
context.ep_broadcast(torch.tensor([1.0]), src_ep_rank=4)
+@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)
@@ -388,22 +412,23 @@ 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,
+def test_vocab_parallel_embedding_reduces_results():
+ reduced = torch.randn(2, 4)
+ context = SimpleNamespace(
+ tp_rank=0,
+ tp_size=2,
+ tp_all_reduce=Mock(return_value=reduced),
)
+ with patch(
+ "sparsevllm.layers.embed_head.get_parallel_context",
+ return_value=context,
+ ):
+ embedding = VocabParallelEmbedding(8, 4)
- actual = ParallelContext._all_reduce(input_tensor, group)
+ output = embedding(torch.tensor([0, 5]))
- assert actual is output_tensor
- provider.run.assert_called_once_with(input_tensor)
+ assert output is reduced
+ context.tp_all_reduce.assert_called_once()
def test_cache_kv_heads_depend_on_tp_not_ep():
diff --git a/tests/test_prefill_attention_provider.py b/tests/test_prefill_attention_provider.py
new file mode 100644
index 00000000..6671893e
--- /dev/null
+++ b/tests/test_prefill_attention_provider.py
@@ -0,0 +1,309 @@
+import gc
+import weakref
+from types import SimpleNamespace
+from unittest.mock import Mock, patch
+
+import pytest
+import torch
+
+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
+
+
+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
+
+
+@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
+ 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)
+ 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 969359a6..4bc8ef04 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",
@@ -1299,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(
@@ -1328,8 +1350,10 @@ 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]
+ runner.eager_static_count = 0
return runner
def make_seq(self, *, prompt_len=100, max_tokens=900, num_tokens=101):
@@ -3745,7 +3769,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 +3799,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 +3821,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_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_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_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 b9ba58ba..99fb48bb 100644
--- a/tests/test_qwen35_hd256_decode_routing.py
+++ b/tests/test_qwen35_hd256_decode_routing.py
@@ -3,10 +3,14 @@
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 (
+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,
)
@@ -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):
@@ -174,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,
@@ -191,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,
@@ -256,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_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_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_rkv_skipkv_methods.py b/tests/test_rkv_skipkv_methods.py
index ecaeac7a..7070d825 100644
--- a/tests/test_rkv_skipkv_methods.py
+++ b/tests/test_rkv_skipkv_methods.py
@@ -4,11 +4,18 @@
import unittest
from unittest.mock import patch
+import numpy as np
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.storage import MlaLatentStorage
from sparsevllm.engine.cache_manager.skipkv import (
SkipKVCacheManager,
SkipKVSentence,
@@ -16,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,
@@ -94,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(
[
@@ -502,9 +526,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,
@@ -532,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
@@ -599,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_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(
diff --git a/tests/test_sgl_fa3.py b/tests/test_sgl_fa3.py
new file mode 100644
index 00000000..adaf1eb3
--- /dev/null
+++ b/tests/test_sgl_fa3.py
@@ -0,0 +1,394 @@
+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,
+)
+
+
+def test_sgl_fa3_support_rejects_missing_package() -> None:
+ with patch("importlib.util.find_spec", return_value=None):
+ assert sgl_fa3_support() == (False, "sglang-kernel is not installed")
+
+
+@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=version),
+ ):
+ supported, reason = sgl_fa3_support()
+
+ assert not supported
+ assert "sglang-kernel>=0.4.5,<0.4.6" in reason
+
+
+def test_sgl_fa3_support_accepts_declared_range() -> None:
+ version = "0.4.5"
+ 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=version),
+ patch("importlib.import_module", return_value=object()),
+ ):
+ supported, reason = sgl_fa3_support()
+
+ assert supported
+ 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.4.5"),
+ 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 sglang-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 sglang-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 sglang-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_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)
+ 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),
+ )
+ 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):
+ 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,
+ )
+ torch.testing.assert_close(packed_output, output)
diff --git a/tests/test_sgl_moe.py b/tests/test_sgl_moe.py
new file mode 100644
index 00000000..82386f19
--- /dev/null
+++ b/tests/test_sgl_moe.py
@@ -0,0 +1,315 @@
+from __future__ import annotations
+
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from sparsevllm.kernels.external.sgl.moe import (
+ sgl_moe_align_block_size,
+ sgl_moe_alignment_support,
+)
+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:
+ with patch("importlib.util.find_spec", return_value=None):
+ assert sgl_moe_alignment_support() == (
+ False,
+ "sglang-kernel is not installed",
+ )
+
+
+@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=version),
+ ):
+ supported, reason = sgl_moe_alignment_support()
+
+ assert not supported
+ assert "sglang-kernel>=0.4.5,<0.4.6" in reason
+
+
+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=version),
+ patch("importlib.import_module", return_value=module),
+ ):
+ supported, reason = sgl_moe_alignment_support()
+
+ assert supported
+ 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.4.5"),
+ 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 sglang-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,
+ )
+ 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
+
+
+@pytest.mark.skipif(
+ not torch.cuda.is_available() or not sgl_moe_alignment_support()[0],
+ 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(
+ 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_alignment_support()[0],
+ 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])
+@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")
+ 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:
+ remote_start = num_local_experts if local_start == 0 else 0
+ topk_ids = torch.randint(
+ remote_start,
+ remote_start + num_local_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,
+ )
+ 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,
+ dtype=torch.bfloat16,
+ device=device,
+ )
+ topk_weights /= topk_weights.sum(dim=-1, keepdim=True)
+ kwargs = {
+ "num_experts": num_experts,
+ "local_expert_start": local_start,
+ }
+
+ 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)
+
+ 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)
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_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_tilelang_mla_kernel.py b/tests/test_tilelang_mla_kernel.py
new file mode 100644
index 00000000..6cec7701
--- /dev/null
+++ b/tests/test_tilelang_mla_kernel.py
@@ -0,0 +1,204 @@
+from __future__ import annotations
+
+import pytest
+import torch
+
+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",
+)
+
+
+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..29d3870f
--- /dev/null
+++ b/tests/test_tilelang_mla_operator.py
@@ -0,0 +1,458 @@
+from __future__ import annotations
+
+import subprocess
+import sys
+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.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,
+ MlaSglFa3Provider,
+ MlaTileLangScoreProvider,
+)
+from sparsevllm.operators.registry import OpResolver
+from sparsevllm.platforms import DeviceCaps, PlatformEnum
+
+
+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",
+ )
+
+
+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)],
+)
+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()
+ assert fa3.call_args.kwargs["num_splits"] == 0
+
+
+@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,
+ )
diff --git a/tests/test_tp_rpc.py b/tests/test_tp_rpc.py
index c5a05448..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()
@@ -443,6 +469,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(
@@ -450,7 +479,7 @@ 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"),
)
with (
@@ -469,6 +498,8 @@ def test_model_runner_exit_drains_graphs_before_barrier():
"sync",
"clear_graphs",
"sync",
+ "close_ops",
+ "sync",
"close_shm",
"barrier",
"unlink_shm",
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 a0c29f5a..840b1c40 100644
--- a/tests/test_triton_moe.py
+++ b/tests/test_triton_moe.py
@@ -5,14 +5,21 @@
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.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():
+ 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(