Compare commits

..
Author SHA1 Message Date
Bugen Zhao 76c973e13c refactor & simplify
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-05 06:39:10 +00:00
Bugen Zhao 53275a22d6 bridge rust tool parser to python
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-05 05:57:20 +00:00
Bugen Zhao 9729e05917 fix minimax m3 tool parser rust unit test
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-05 05:52:34 +00:00
Bugen ZhaoandGitHub be37de73a9 Python ame tool/reasoning parser (#9)
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
2026-06-04 10:33:55 +08:00
Hongxia YangandGitHub e633514f50 [ROCm] minimax-m3 initial enablement on ROCm (#25)
* initial enablement on ROCm

Signed-off-by: Hongxia Yang <hongxia.yang@amd.com>

* fuse gemma_rmsnorm, swiglu-oai Triton kernels and tune MXFP8 GEMMs for ROCm

Replace per-element PyTorch fallbacks with single-pass Triton kernels for
the AMD/ROCm MiniMax-M3 path:

- gemma_rmsnorm / gemma_fused_add_rmsnorm: fp32 normalize+scale in one
  kernel, handles non-contiguous qkv.split views (q_norm/k_norm).
- swiglu_oai_split: fused gate*sigmoid(alpha*gate)*(up+beta) with optional
  clamp, used in both dense MLP and native MXFP8 MoE.
- MXFP8 GEMM tuning: set num_warps=8 for both the linear and MoE
  dot_scaled kernels; use hidden_states.dtype instead of hardcoded bf16.

Signed-off-by: Hongxia Yang <hongxia.yang@amd.com>

* add unit tests for ROCm fused kernels (rmsnorm, swiglu, mxfp8)

Signed-off-by: Hongxia Yang <hongxia.yang@amd.com>

* [ROCm] Add NATIVE_MXFP8 backend, use supports_mx() for arch gating, fix lint

Signed-off-by: Hongxia Yang <hongxia.yang@amd.com>

* add adaptive tile to swiglu-oai kernal

Signed-off-by: Hongxia Yang <hongxia.yang@amd.com>

* [ROCm] mimimax-m3: enable vision-language support, mirroring nvda path

Signed-off-by: Hongxia Yang <hongxia.yang@amd.com>

* [ROCm] minimax-m3: review feedback on SWIGLUOAI_UNINTERLEAVE and tidy MXFP8 MoE oracle

Signed-off-by: Hongxia Yang <hongxia.yang@amd.com>

* revert changes to nvda mxfp8 oracle to minimize changes

Signed-off-by: Hongxia Yang <hongxia.yang@amd.com>

* minor

Signed-off-by: Hongxia Yang <hongxia.yang@amd.com>

---------

Signed-off-by: Hongxia Yang <hongxia.yang@amd.com>
2026-06-03 22:04:27 -04:00
Jee Jee LiandGitHub 6647a1a88b Reduce dtype conversions for AR fusion (#40)
* Init

Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>

* move

Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>

---------

Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
2026-06-04 00:43:45 +08:00
Jee Jee LiandGitHub 56afa45bf7 Fuse qknorm+rope+insert (#20)
* Init

Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>

* Delete comments

Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>

* NHD layout

Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>

* torch stable abi

Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>

* torch stable abi

Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>

* FMT

Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>

* FMT

Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>

---------

Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
2026-06-03 23:10:47 +08:00
Jiangyun ZhuandGitHub 41cf8bbd60 enable breakable cudagraph by default (#39)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
2026-06-03 20:17:46 +08:00
Bugen ZhaoandGitHub 427b2793f0 Rust frontend: renderer fixes, tool parser, reasoning parser (#8)
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
Fix several hf chat template rendering issues (#44311)
2026-06-03 15:12:19 +08:00
Thien TranandGitHub f4ab9994b5 Simple improvement to Indexer score prefill (#37)
* improve

* remove masked load for K

* add comment
2026-06-03 14:59:15 +08:00
8c9b156eca HND KV cache layout (#29)
* HND layout

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>

* refactor: integrate sparse KV cache into stride-order layout contract

Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>

---------

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>
2026-06-03 08:59:46 +08:00
738af995a7 [MiniMax-M3] Enable EAGLE3 speculative decoding (#24)
* [MiniMax-M3] Enable EAGLE3 speculative decoding

Signed-off-by: zixi-qi <zixi@inferact.ai>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* [MiniMax-M3] Document that PP is unsupported for EAGLE3 aux indices

Signed-off-by: zixi-qi <zixi@inferact.ai>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

---------

Signed-off-by: zixi-qi <zixi@inferact.ai>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-02 13:54:29 -04:00
Jee Jee LiandGitHub 293c3895e2 Done (#32)
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
2026-06-02 20:43:48 +08:00
e7fef86e50 [CI/Build] Bump flashinfer to v0.6.12 (#31)
(cherry picked from commit ef2acc7fb7)

Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com>
Co-authored-by: Vadim Gimpelson <vadim.gimpelson@gmail.com>
2026-06-02 16:13:25 +08:00
Thien TranandGitHub 9cb8ed5008 Remove torch.full() in indexer kernels (#30)
Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>
2026-06-02 15:26:17 +08:00
c5905f7760 [MiniMax M3] Vision-language (ViT) support + flashinfer_cudnn & DP-encoder features (#16)
* [Model] Add MiniMax M3 text backbone skeleton + SwiGLU-OAI clamp activation

Port the MiniMax M3 (text backbone) into vLLM's custom model layout:

- Add MiniMaxM3SparseForCausalLM under vllm/models/minimax_m3/nvidia with the
  decoder/model/causal-LM wiring; attention and MoE bodies plus weight loading
  are left as stubs. Dense MiniMaxM3MLP is fully ported.
- Add MiniMaxM3SparseForConditionalGeneration as a minimal LM-routing wrapper
  (KimiK25-style init_vllm_registered_model on text_config) and register both
  architectures.
- Add MiniMaxM3Config (model_type minimax_m3_vl) wrapping MiniMaxM3TextConfig
  so config.get_text_config() extracts the backbone; register in the config
  registries.

Generalize silu_and_mul_with_clamp to SwiGLU-OAI:

- Add alpha (scales the activation's sigmoid) and beta (added to the
  non-activated half) to the CUDA kernel, ops.h, and torch_bindings schema.
  Defaults alpha=1.0, beta=0.0 are bitwise-identical to the previous
  silu(gate)*up, so existing callers (DeepSeek V4) are unaffected.
- SiluAndMulWithClamp(alpha, beta) used by MiniMaxM3MLP with alpha=swiglu_alpha,
  beta=1.0, matching the reference gate*sigmoid(alpha*gate)*(up+1).

AI assistance (Claude) was used for this change.

Signed-off-by: Yongye Zhu <yongye@inferact.ai>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>

* [Model] MiniMax M3: implement MoE block + weight-name mapping

Implement the sigmoid-routed MoE block for MiniMax M3 and map module
names to the checkpoint so weight loading works for the ported modules.

MoE block (MiniMaxM3MoE):
- fp32 router via GateLinear (bf16 activations upcast to fp32; fp32
  weights and logits), matching minimax_m2/sglang precision.
- FusedTopKBiasRouter routing (scoring_func from config, sigmoid +
  e_score_correction_bias + renormalize), verified to match sglang's
  TopK (select-with-bias, weight-without-bias, routed_scaling on output).
- swigluoai activation (from config.hidden_act) + swiglu_limit; shared
  expert fused into FusedMoE so the shared partial is reduced with the
  routed output.

Weight loading:
- Name the MoE submodule `block_sparse_moe` (dense stays `mlp`) to match
  the checkpoint; decoder forward selects per layer.
- load_weights handles gate_up fusion (dense MLP + shared experts) and
  expert w1/w2/w3 -> w13/w2 fusion; wrappers delegate via
  AutoWeightsLoader, skipping vision/mm/mtp. Not-yet-ported modules
  (attention) are skipped until they land.

The expert GEMM/activation kernel correctness and attention/weight
loading for the remaining modules are not part of this change.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>

* [MiniMax M3] Port attention modules and wire up MXFP8 checkpoint loading

Port the weight-bearing dense (MiniMaxM3Attention) and sparse
(MiniMaxM3SparseAttention) attention modules so the checkpoint's
self_attn.* tensors map onto real params (forward still stubbed;
this targets weight loading). Add qkv stacked mapping and the
weight_scale_inv -> weight_scale remap in load_weights.

Load MiniMax-style MXFP8 checkpoints (quant_method: "mxfp8" +
ignored_layers) via the ModelOpt MXFP8 config: register "mxfp8" in
method_to_config and normalize the minimal checkpoint schema to the
ModelOpt schema in ModelOptMxFp8Config.from_config (same on-disk
format). Use setdefault for online shorthands so the checkpoint
config wins over the "mxfp8" online shorthand.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Yongye Zhu <yongye@inferact.ai>

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>

* [MiniMax M3] Add DeepGEMM MXFP8 MoE backend with swigluoai support

Extend DeepGemmExperts to support MXFP8 activations (FP8 e4m3 + UE8M0
1x32 block scales) via the grouped GEMM with recipe (1, 32), reusing the
oracle/fp8 weight-conversion path. Generalize
deepgemm_post_process_fp8_weight_block to derive the transform recipe
from the block shape ((1, 1, 32) for MXFP8) and accept uint8 E8M0 scales.

Unify the fused gated-activation+quant triton kernels around
y = (up + beta) * gate * sigmoid(alpha * gate): silu is alpha=1, beta=0
(bit-identical to before); swigluoai uses alpha/beta from config. Thread
gemm1_alpha/gemm1_beta from the FusedMoE layer through the MXFP8 quant
config into the kernels, and add swiglu_alpha/swiglu_beta to the layer
and MiniMax M3 config/model (beta sourced from config, not hardcoded).

Wire Fp8MoeBackend.DEEPGEMM into the MXFP8 oracle (selectable via
--moe-backend deep_gemm), resolving directly to DeepGemmExperts (the
Triton fallback cannot handle the 1x32 scheme). Advertise SWIGLUOAI in
_supports_activation so swigluoai selects DeepGEMM rather than falling
through to another backend; gate the MXFP8 scheme to Blackwell (SM100).

Verified on GB200: packed-kernel parity (silu defaults unchanged,
swigluoai matches torch ref), (1,32) weight-prep transform, and a TP=4
launch selecting the DEEPGEMM MXFP8 backend with full weight load (the
run then stops at the still-stubbed attention forward, as expected).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Yongye Zhu <yongye@inferact.ai>

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>

* [MiniMax M3] Implement dense attention forward

Implement MiniMaxM3Attention.forward (dense path): qkv projection, split,
per-head QK norm (GemmaRMSNorm, qk_norm_type="per_head"), partial RoPE,
attention, and output projection. Mirrors the sglang reference dense path
and vLLM's canonical per-head-norm convention. attention_output_gate is
False for M3, so the gate branch is omitted.

The sparse attention forward (index branch) remains stubbed.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Yongye Zhu <yongye@inferact.ai>

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>

* [MiniMax M3] Drop dead index value/output projection branch

For M3, sparse_disable_index_value matches sparse_attention_freq exactly
([0,0,0,1,...,1]): the only layers with the flag unset (0-2) are the
non-sparse layers built as MiniMaxM3Attention. Every layer that constructs
MiniMaxM3SparseAttention therefore always disables the index value/output
projections, so index_{v,o}_proj are never created.

Remove the unreachable else branch, the disable_index_value parameter and
field, and the now-unused _disable_index_value_layer_ids helper.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Yongye Zhu <yongye@inferact.ai>

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>

* [MiniMax M3] Add sparse-attention backend + merged attention layer

Scaffold the lightning-indexer sparse-attention path:

- MiniMaxM3SparseBackend (registered as MINIMAX_M3_SPARSE): block-sparse GQA
  backend; get_kv_cache_shape serves both the main K/V cache and the
  single-vector index-key side cache.
- MiniMaxM3IndexerCache: side KV cache for per-token index keys, key-only so it
  uses a single-vector MLAAttentionSpec rather than a K+V FullAttentionSpec.
- MiniMaxM3SparseMetadata (+ prefill/decode sub-metadata) and its builder,
  splitting the batch via split_decodes_and_prefills.
- MiniMaxM3SparseImpl: subclasses AttentionImplBase so it can take a custom
  forward(query, index_query, kv_cache, index_kv_cache); no alibi / sliding
  window / logits soft cap. forward is a stub pending the kernel port.

MiniMaxM3SparseAttention is merged into a single AttentionLayerBase: it owns the
projections, per-head QK norm and RoPE, binds the backend + impl, registers the
main K/V cache, and holds the index cache. Its forward computes q/k/v and the
index q/k, pre-inserts K/V and index-K into their caches, then calls the sparse
impl with only the queries.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Yongye Zhu <yongye@inferact.ai>

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>

* [MiniMax M3] Implement sparse-attention forward + MXFP8 DeepGEMM MoE e2e

Port the MiniMax M3 lightning-indexer sparse attention to a Triton backend and
fix the MXFP8 DeepGEMM MoE path so the model runs end to end.

Attention (vllm/v1/attention/ops/minimax_m3_sparse_ops.py + backend):
- Triton kernels (paged, page == sparse block == 128): index block-score +
  bitonic top-k, and GQA block-sparse flash attention over the selected blocks.
- MiniMaxM3SparseImpl.forward: decode-first split, dispatching the same kernels
  per phase (a decode token is a 1-token prefill). Index and main caches use
  separate block tables.
- Dedicated MiniMaxM3IndexerBackend for the key-only index cache so the main
  GQA cache (num_kv_heads==1 at TP>=4) is not mistaken for the index layout.
- get_supported_kernel_block_sizes()==[128] (one sparse block per KV page).

MXFP8 DeepGEMM MoE:
- Prepare-phase activation quant emits float32 per-(1,32) group scales for the
  DeepGEMM backend (use_deep_gemm_packed_mxfp8 on the quant config), matching
  the FP8 128-block path with group=32.
- deepgemm_moe_permute / ep_scatter take a block_size so the activation-scale
  group (32) is honored through the expert permute.
- workspace_shapes uses the contiguous-layout M alignment (not block_shape[0],
  which is 1 for MXFP8 and under-sized the workspace).

GSM8K (5-shot, TP=4) flexible-extract 0.921 / strict 0.919.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Yongye Zhu <yongye@inferact.ai>

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>

* draft vl impl

Signed-off-by: Isotr0py <Isotr0py@outlook.com>

* try to load

Signed-off-by: Isotr0py <Isotr0py@outlook.com>

* don't rename o_proj

Signed-off-by: Isotr0py <Isotr0py@outlook.com>

* fix vit loading

Signed-off-by: Isotr0py <Isotr0py@outlook.com>

* ooops

Signed-off-by: Isotr0py <Isotr0py@outlook.com>

* ooops

Signed-off-by: Isotr0py <Isotr0py@outlook.com>

* [MiniMax M3] Enable decode CUDA graphs + dedicated split-K decode kernels (#7)

Two changes to the sparse-attention backend:

1. Full decode CUDA-graph support. The metadata builder now declares
   AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE and precomputes all
   per-step kernel args (cu_seqlens_q, prefix_lens, max_query_len,
   num_actual_tokens) in build(), removing the .item() host sync and the
   per-step torch.zeros/cumsum/diff from the impl forward. Derived decode
   prefix lengths are written into a persistent buffer so the captured
   graph reads stable addresses across replays.

2. Dedicated split-K decode kernels (mirroring the sglang reference)
   instead of reusing the prefill kernels with BLOCK_SIZE_Q=1, which left
   the GPU idle at decode (one query token per request). The index score
   now splits over seq blocks and the GQA attention splits over the
   selected top-k blocks with an LSE merge (flash-decoding). Chunk counts
   depend only on shape constants, so the grid is fixed within a CUDA
   graph.

Verified: a parity test against the prior (GSM8K 92.1) prefill-as-decode
path matches exactly on top-k selection and on attention output (bf16
noise) across seq lengths 128-2048.

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>

* [MiniMax M3 VL] Multimodal (ViT) support: vendored processor, ViT fixes, encoder features

Make MiniMaxAI/Minimax-M3-preview serve as a VL model:
- registry: move ...ForConditionalGeneration to _MULTIMODAL_MODELS
- vendor the HF processor (image/video/composite) so no --trust-remote-code,
  constructed directly in get_hf_processor; Qwen-style smart_resize
- vision_tower: disable post_layernorm (matches reference), fp32 RoPE,
  backend-aware encoder metadata enabling flashinfer_cudnn ViT
- model: supports_encoder_tp_data + fix --mm-encoder-tp-mode data DP branch
- mm_preprocess: cap dummy video frames; smart_resize token counting

AI-assisted (Claude Code); WIP, pending human review + test re-run.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* [MiniMax M3 VL] Align video timestamps with the MiniMax reference

Request video metadata (fps + sampled frame indices) via
MultiModalDataParser(video_needs_metadata=True) and forward it as
VideoMetadata so the processor emits per-frame "]<]X.X seconds[>[" markers.
_get_prompt_updates reconstructs the same markers from the metadata (using the
HF formula frames_indices[frame*temporal_patch_size]/fps) so the prompt
replacement stays byte-aligned with the processor output. Falls back to no
timestamps when metadata is absent (dummy/profiling videos), keeping both
paths consistent.

Verified: processor emits the expected timestamps; the piecewise replacement
exactly matches the tokenized video region; server video requests succeed
(no placeholder mismatch) and image inference is unaffected.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* m3 video loader and processor cleanup

Signed-off-by: Isotr0py <Isotr0py@outlook.com>

* Fix pre-commit lint/type errors in MiniMax M3 VL files

The new MiniMax M3 VL files live outside vllm/model_executor/models (which
mypy excludes), so they are linted and type-checked in CI. Make all hooks
pass:

- ruff: reorder default_weight_loader import (isort); drop trailing
  whitespace in vision_tower.py.
- typos: allowlist `tpos` (temporal position id, parallels the existing
  hpos/wpos vision-RoPE naming).
- mypy:
  - annotate round/ceil/floor_by_factor as `int | float` (matches
    ernie45_vl); they are called with float values.
  - `# type: ignore[call-arg]` on the ImagesKwargs/VideosKwargs/
    ProcessingKwargs `total=False` subclasses (matches ovis/isaac/etc.).
  - cast dummy-option overrides to ImageDummyOptions/VideoDummyOptions and
    the videos mm_data value to `list` before iterating.
  - assert multimodal_config is not None before reading mm_encoder_tp_mode.
  - skip None results when building mm_input_by_modality so the strict
    `dict[str, dict]` annotation holds and no None reaches the embedders.

Verified locally: ruff-check, ruff-format, typos, check-spdx-header,
check-root-lazy-imports, and mypy (3.10/3.11/3.12/3.13) all pass on the
changed files.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Roger Wang <hey@rogerw.io>

---------

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
Signed-off-by: Isotr0py <Isotr0py@outlook.com>
Signed-off-by: Roger Wang <hey@rogerw.io>
Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>
Co-authored-by: Isotr0py <Isotr0py@outlook.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-01 22:08:42 -07:00
22c6542fa7 [FlashInfer] Support page size >= 128 for trtllm-gen GQA/MQA (#26)
Add power-of-2 page sizes >= 128 (128/256/512/1024) to the FlashInfer
backend's supported kernel block sizes. These are served only by the
trtllm-gen dynamic kernel, which requires Blackwell + GQA/MQA
(num_qo_heads // num_kv_heads > 1), not MHA.

Fail fast in FlashInferMetadataBuilder.__init__ when page_size >= 128 is
requested without a usable trtllm-gen path (trtllm disabled, not
Blackwell, or MHA) instead of silently hitting the native wrappers, and
force the prefill path to trtllm so it does not fall back to FA2.

Regenerate docs/design/attention_backends.md for the new block sizes.

AI assistance (Claude) was used for this change.

Co-authored-by: Claude

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
2026-06-02 00:30:09 -04:00
Yongye ZhuandGitHub 483bda03a7 [MiniMax-M3] Fuse indexer q/k into qkv_proj; support TP > kv heads (#23)
The sparse-attention layers ran three separate projections off the same
hidden_states (qkv_proj + index_q_proj + index_k_proj) and broke for
tensor-parallel sizes greater than num_key_value_heads (4): num_idx_heads
collapsed to 0 and the index_q ColumnParallelLinear could not split 4 heads
across >4 ranks.

Add MinimaxM3QKVParallelLinearWithIndexer (in layers/linear.py), a
QKVParallelLinear subclass that emits [q | k | v | index_q | index_k] from a
single column-parallel GEMM. Since sparse_num_index_heads == num_key_value_heads
and index_head_dim == head_dim, index_q shards exactly like K/V and rides the
existing num_kv_head_replicas path, so TP > kv-heads works (replication) and the
num_idx_heads == num_kv_heads indexer invariant holds at every TP. index_k is a
single shared head replicated to every rank. MXFP8 weight_scale shards via the
same per-output-row offsets (block_shape [1, 32]). The class is M3-specific; it
lives in linear.py only to reuse QKVParallelLinear's sharding/loading machinery.

Wire MiniMaxM3SparseAttention to the fused linear (single projection + 5-way
split) and fold index_q/index_k_proj weights into qkv_proj in the loader; these
mapping entries never match dense layers, which keep plain QKVParallelLinear.

Verified TP=4 on GB200: weights load clean, GSM8K 5-shot exact_match
flexible-extract 0.9189, strict-match 0.9166.

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
2026-06-01 22:39:20 -04:00
92b3a243d5 [MiniMax-M3] Add SWIGLUOAI_UNINTERLEAVE activation and fix MXFP8 Marlin MoE (#22)
MiniMax-M3 loads w13 packed ([all gates; all ups]) via
MergedColumnParallelLinear, but the existing SWIGLUOAI activation assumes
gpt-oss-style interleaved w13. Add a SWIGLUOAI_UNINTERLEAVE variant that runs
the same SwiGLU-OAI math on packed w13 via silu_and_mul_with_clamp, and wire it
through Marlin and deepgemm.

- activation: add SWIGLUOAI_UNINTERLEAVE; extend apply_moe_activation with
  clamp_limit/alpha/beta and fuse the SILU+clamp (swiglu_limit_func) path into
  silu_and_mul_with_clamp.
- marlin_moe: route the new activation, plumb gemm1_alpha/beta from the quant
  config.
- modular_kernel: forward the gate params through activation().
- deep_gemm_moe: switch its packed gated path to SWIGLUOAI_UNINTERLEAVE.
- oracle/fp8 + config: fp8_w8a16_moe_quant_config (MARLIN/CPU branch) was
  dropping swiglu_limit/alpha/beta, leaving gemm1_clamp_limit=None on the MXFP8
  Marlin path; forward them.
- minimax_m3 model: emit "swigluoai_uninterleave".
- minimax_m3 config: swiglu_beta default 1.0 (the (up+1) OAI bias; the
  checkpoint omits it). With beta=0 the +1 term is dropped on every FFN/MoE
  layer.

gsm8k (MiniMax-M3-preview, TP=4, Hopper/Marlin MXFP8): 0.0008 -> 0.92.

Signed-off-by: yongye <yongye@inferact.ai>
Co-authored-by: yongye <yongye@inferact.ai>
2026-06-01 20:04:13 -04:00
14682903c8 Improve silu_mul_quant_fp8 (#17)
* Optimize packed SiLU FP8 quant tile

Record MiniMax M3 workload benchmark results.

Co-authored-by: OpenAI Codex <codex@openai.com>

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>

* Record power-of-four SiLU FP8 workload sweep

Co-authored-by: OpenAI Codex <codex@openai.com>

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>

* Compare power-of-four SiLU FP8 sweep

Co-authored-by: OpenAI Codex <codex@openai.com>

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>

* Use one-row CTAs for packed SiLU FP8 quant

Record CUPTI benchmark comparison for the one-row CTA variant.

Co-authored-by: OpenAI Codex <codex@openai.com>

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>

* Use persistent row CTAs for packed SiLU FP8 quant

Cap row CTAs and stride across M to avoid oversized grids while retaining large-shape throughput.

Co-authored-by: OpenAI Codex <codex@openai.com>

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>

* Use flat persistent loads for packed SiLU FP8 quant

Record latest CUPTI progression and rejected cache/load experiments.

Co-authored-by: OpenAI Codex <codex@openai.com>

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>

* final kernel

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>

* Remove SiLU FP8 worklog from tracking

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>

* rename

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>

---------

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>
2026-06-02 03:42:36 +08:00
cbfaaeceeb [MiniMax-M3] Fuse MXFP8 UE8M0 scale packing into ep_scatter (#21)
Remove the `use_deep_gemm_packed_mxfp8` flag and the special MXFP8
quantization branch in `moe_kernel_quantize_input`. MXFP8 activations are
now always quantized to the plain non-swizzled (M, K/32) uint8 UE8M0 scale
layout, and the pack into DeepGEMM's consumed scale layout (int32, MN-major,
TMA-aligned, 4 UE8M0 per int32) is fused directly into the expert-permute
scatter instead of relying on the GEMM's internal repack.

- utils.py: drop the flag + branch; mxfp8 always uses _mxfp8_e4m3_quantize
  (non-swizzled).
- config.py / no_dp_ep.py / oracle/fp8.py: remove the flag plumbing.
- deep_gemm_utils.py: add a PACK_UE8M0 path to _fwd_kernel_ep_scatter_2 that
  concatenates 4 UE8M0 bytes per int32 and stores MN-major; deepgemm_moe_permute
  allocates the TMA-aligned int32 buffer for the uint8 path. Float32 (FP8/FP4)
  scales keep the row-major path unchanged.

mm1 now feeds the grouped GEMM pre-packed int32 scales with recipe_a=(1,32),
matching what the mm2 activation-quant path already does (validate-only
transform).

AI assistance (Claude Code) was used for this change.

Tests run on GB200 (SM100):
- Numerical unit test: fused packed scatter matches a torch reference for
  data placement and byte-packing; output layout/stride is identical to
  per_token_group_quant_fp8_packed_for_deepgemm.
- gsm8k 5-shot, MiniMax-M3-preview:
    TP=4 monolithic : exact_match 0.9249 +/- 0.0073
    DP=4 + EP       : exact_match 0.9325 +/- 0.0069


Co-authored-by: Claude

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-01 15:06:21 -04:00
cf6c0d2518 Add AR GemmaRMSNorm fusion (#11)
* [Model] Add MiniMax M3 text backbone skeleton + SwiGLU-OAI clamp activation

Port the MiniMax M3 (text backbone) into vLLM's custom model layout:

- Add MiniMaxM3SparseForCausalLM under vllm/models/minimax_m3/nvidia with the
  decoder/model/causal-LM wiring; attention and MoE bodies plus weight loading
  are left as stubs. Dense MiniMaxM3MLP is fully ported.
- Add MiniMaxM3SparseForConditionalGeneration as a minimal LM-routing wrapper
  (KimiK25-style init_vllm_registered_model on text_config) and register both
  architectures.
- Add MiniMaxM3Config (model_type minimax_m3_vl) wrapping MiniMaxM3TextConfig
  so config.get_text_config() extracts the backbone; register in the config
  registries.

Generalize silu_and_mul_with_clamp to SwiGLU-OAI:

- Add alpha (scales the activation's sigmoid) and beta (added to the
  non-activated half) to the CUDA kernel, ops.h, and torch_bindings schema.
  Defaults alpha=1.0, beta=0.0 are bitwise-identical to the previous
  silu(gate)*up, so existing callers (DeepSeek V4) are unaffected.
- SiluAndMulWithClamp(alpha, beta) used by MiniMaxM3MLP with alpha=swiglu_alpha,
  beta=1.0, matching the reference gate*sigmoid(alpha*gate)*(up+1).

AI assistance (Claude) was used for this change.

Signed-off-by: Yongye Zhu <yongye@inferact.ai>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>

* [Model] MiniMax M3: implement MoE block + weight-name mapping

Implement the sigmoid-routed MoE block for MiniMax M3 and map module
names to the checkpoint so weight loading works for the ported modules.

MoE block (MiniMaxM3MoE):
- fp32 router via GateLinear (bf16 activations upcast to fp32; fp32
  weights and logits), matching minimax_m2/sglang precision.
- FusedTopKBiasRouter routing (scoring_func from config, sigmoid +
  e_score_correction_bias + renormalize), verified to match sglang's
  TopK (select-with-bias, weight-without-bias, routed_scaling on output).
- swigluoai activation (from config.hidden_act) + swiglu_limit; shared
  expert fused into FusedMoE so the shared partial is reduced with the
  routed output.

Weight loading:
- Name the MoE submodule `block_sparse_moe` (dense stays `mlp`) to match
  the checkpoint; decoder forward selects per layer.
- load_weights handles gate_up fusion (dense MLP + shared experts) and
  expert w1/w2/w3 -> w13/w2 fusion; wrappers delegate via
  AutoWeightsLoader, skipping vision/mm/mtp. Not-yet-ported modules
  (attention) are skipped until they land.

The expert GEMM/activation kernel correctness and attention/weight
loading for the remaining modules are not part of this change.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>

* [MiniMax M3] Port attention modules and wire up MXFP8 checkpoint loading

Port the weight-bearing dense (MiniMaxM3Attention) and sparse
(MiniMaxM3SparseAttention) attention modules so the checkpoint's
self_attn.* tensors map onto real params (forward still stubbed;
this targets weight loading). Add qkv stacked mapping and the
weight_scale_inv -> weight_scale remap in load_weights.

Load MiniMax-style MXFP8 checkpoints (quant_method: "mxfp8" +
ignored_layers) via the ModelOpt MXFP8 config: register "mxfp8" in
method_to_config and normalize the minimal checkpoint schema to the
ModelOpt schema in ModelOptMxFp8Config.from_config (same on-disk
format). Use setdefault for online shorthands so the checkpoint
config wins over the "mxfp8" online shorthand.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Yongye Zhu <yongye@inferact.ai>

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>

* [MiniMax M3] Add DeepGEMM MXFP8 MoE backend with swigluoai support

Extend DeepGemmExperts to support MXFP8 activations (FP8 e4m3 + UE8M0
1x32 block scales) via the grouped GEMM with recipe (1, 32), reusing the
oracle/fp8 weight-conversion path. Generalize
deepgemm_post_process_fp8_weight_block to derive the transform recipe
from the block shape ((1, 1, 32) for MXFP8) and accept uint8 E8M0 scales.

Unify the fused gated-activation+quant triton kernels around
y = (up + beta) * gate * sigmoid(alpha * gate): silu is alpha=1, beta=0
(bit-identical to before); swigluoai uses alpha/beta from config. Thread
gemm1_alpha/gemm1_beta from the FusedMoE layer through the MXFP8 quant
config into the kernels, and add swiglu_alpha/swiglu_beta to the layer
and MiniMax M3 config/model (beta sourced from config, not hardcoded).

Wire Fp8MoeBackend.DEEPGEMM into the MXFP8 oracle (selectable via
--moe-backend deep_gemm), resolving directly to DeepGemmExperts (the
Triton fallback cannot handle the 1x32 scheme). Advertise SWIGLUOAI in
_supports_activation so swigluoai selects DeepGEMM rather than falling
through to another backend; gate the MXFP8 scheme to Blackwell (SM100).

Verified on GB200: packed-kernel parity (silu defaults unchanged,
swigluoai matches torch ref), (1,32) weight-prep transform, and a TP=4
launch selecting the DEEPGEMM MXFP8 backend with full weight load (the
run then stops at the still-stubbed attention forward, as expected).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Yongye Zhu <yongye@inferact.ai>

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>

* [MiniMax M3] Implement dense attention forward

Implement MiniMaxM3Attention.forward (dense path): qkv projection, split,
per-head QK norm (GemmaRMSNorm, qk_norm_type="per_head"), partial RoPE,
attention, and output projection. Mirrors the sglang reference dense path
and vLLM's canonical per-head-norm convention. attention_output_gate is
False for M3, so the gate branch is omitted.

The sparse attention forward (index branch) remains stubbed.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Yongye Zhu <yongye@inferact.ai>

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>

* [MiniMax M3] Drop dead index value/output projection branch

For M3, sparse_disable_index_value matches sparse_attention_freq exactly
([0,0,0,1,...,1]): the only layers with the flag unset (0-2) are the
non-sparse layers built as MiniMaxM3Attention. Every layer that constructs
MiniMaxM3SparseAttention therefore always disables the index value/output
projections, so index_{v,o}_proj are never created.

Remove the unreachable else branch, the disable_index_value parameter and
field, and the now-unused _disable_index_value_layer_ids helper.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Yongye Zhu <yongye@inferact.ai>

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>

* [MiniMax M3] Add sparse-attention backend + merged attention layer

Scaffold the lightning-indexer sparse-attention path:

- MiniMaxM3SparseBackend (registered as MINIMAX_M3_SPARSE): block-sparse GQA
  backend; get_kv_cache_shape serves both the main K/V cache and the
  single-vector index-key side cache.
- MiniMaxM3IndexerCache: side KV cache for per-token index keys, key-only so it
  uses a single-vector MLAAttentionSpec rather than a K+V FullAttentionSpec.
- MiniMaxM3SparseMetadata (+ prefill/decode sub-metadata) and its builder,
  splitting the batch via split_decodes_and_prefills.
- MiniMaxM3SparseImpl: subclasses AttentionImplBase so it can take a custom
  forward(query, index_query, kv_cache, index_kv_cache); no alibi / sliding
  window / logits soft cap. forward is a stub pending the kernel port.

MiniMaxM3SparseAttention is merged into a single AttentionLayerBase: it owns the
projections, per-head QK norm and RoPE, binds the backend + impl, registers the
main K/V cache, and holds the index cache. Its forward computes q/k/v and the
index q/k, pre-inserts K/V and index-K into their caches, then calls the sparse
impl with only the queries.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Yongye Zhu <yongye@inferact.ai>

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>

* [MiniMax M3] Implement sparse-attention forward + MXFP8 DeepGEMM MoE e2e

Port the MiniMax M3 lightning-indexer sparse attention to a Triton backend and
fix the MXFP8 DeepGEMM MoE path so the model runs end to end.

Attention (vllm/v1/attention/ops/minimax_m3_sparse_ops.py + backend):
- Triton kernels (paged, page == sparse block == 128): index block-score +
  bitonic top-k, and GQA block-sparse flash attention over the selected blocks.
- MiniMaxM3SparseImpl.forward: decode-first split, dispatching the same kernels
  per phase (a decode token is a 1-token prefill). Index and main caches use
  separate block tables.
- Dedicated MiniMaxM3IndexerBackend for the key-only index cache so the main
  GQA cache (num_kv_heads==1 at TP>=4) is not mistaken for the index layout.
- get_supported_kernel_block_sizes()==[128] (one sparse block per KV page).

MXFP8 DeepGEMM MoE:
- Prepare-phase activation quant emits float32 per-(1,32) group scales for the
  DeepGEMM backend (use_deep_gemm_packed_mxfp8 on the quant config), matching
  the FP8 128-block path with group=32.
- deepgemm_moe_permute / ep_scatter take a block_size so the activation-scale
  group (32) is honored through the expert permute.
- workspace_shapes uses the contiguous-layout M alignment (not block_shape[0],
  which is 1 for MXFP8 and under-sized the workspace).

GSM8K (5-shot, TP=4) flexible-extract 0.921 / strict 0.919.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Yongye Zhu <yongye@inferact.ai>

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>

* [MiniMax M3] Enable decode CUDA graphs + dedicated split-K decode kernels (#7)

Two changes to the sparse-attention backend:

1. Full decode CUDA-graph support. The metadata builder now declares
   AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE and precomputes all
   per-step kernel args (cu_seqlens_q, prefix_lens, max_query_len,
   num_actual_tokens) in build(), removing the .item() host sync and the
   per-step torch.zeros/cumsum/diff from the impl forward. Derived decode
   prefix lengths are written into a persistent buffer so the captured
   graph reads stable addresses across replays.

2. Dedicated split-K decode kernels (mirroring the sglang reference)
   instead of reusing the prefill kernels with BLOCK_SIZE_Q=1, which left
   the GPU idle at decode (one query token per request). The index score
   now splits over seq blocks and the GQA attention splits over the
   selected top-k blocks with an LSE merge (flash-decoding). Chunk counts
   depend only on shape constants, so the grid is fixed within a CUDA
   graph.

Verified: a parity test against the prior (GSM8K 92.1) prefill-as-decode
path matches exactly on top-k selection and on attention output (bf16
noise) across seq lengths 128-2048.

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>

* Init

Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>

* Move

Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>

* FIX

Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>

* Addresss conflict

Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>

---------

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>
2026-06-01 14:14:39 -04:00
Jee Jee LiandGitHub 1b563d1134 Implement MiniMAXGemmaRMSNorm (#19)
* Init

Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>

* Delete comments

Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>

---------

Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
2026-06-02 01:06:48 +08:00
Thien TranandGitHub cbdfa83c84 remove unnecessary padding (#18)
Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>
2026-06-01 12:14:52 -04:00
Thien TranandGitHub 12846bbf88 use FlashInfer RMSNorm (#15)
* use FlashInfer RMSNorm

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>

* restore original GemmaRMSNorm module

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>

---------

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>
2026-06-01 11:21:08 +08:00
f074bd6cef [MiniMax-M3] Group full-attn + MLA-indexer KV in one group; NIXL non-MLA heterogeneous support (#14)
* [MiniMax-M3] Group full-attn + MLA-indexer KV in one group; NIXL non-MLA heterogeneous support

MiniMax-M3 registers full-attention K/V layers plus a smaller MLA "lightning
indexer" cache. These share the token-slot invariant, so they belong in one
UniformTypeKVCacheSpecs group (as DeepSeek-3.2 does for its all-MLA main+indexer),
keeping each layer's tensor at its native page size.

- kv_cache_interface.is_uniform_type: accept a mix of FullAttentionSpec +
  MLAAttentionSpec with the same block_size in one group.
- nixl/worker.register_kv_caches: drop the non-MLA "all tensors same size"
  assert; per-layer sizes are carried in block_len_per_layer and the descriptor
  builders already index them per layer (num_descs uses a common num_blocks).
- nixl/worker.add_remote_agent: for the heterogeneous non-MLA case, validate each
  layer's block_len against the remote (like the MLA path) and raise
  NotImplementedError for tp_ratio != 1 (head-splitting assumes uniform size).

Verified: single-node M3 TP4 (1 group, registration OK with heterogeneous
tensors, APC works, GSM8K 0.96) and 1P1D Slurm PD via NixlConnector (both
engines 1 group, NIXL KV transfers succeed, 8/8 bench requests, 0 failed).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Address comments

Signed-off-by: Dao Le <daole@inferact.ai>

---------

Signed-off-by: Dao Le <daole@inferact.ai>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-05-31 20:16:25 -07:00
Thien TranandGitHub 579d7b3705 Prefill sparse attention CuteDSL kernel (#10)
* Add MiniMax M3 CuteDSL sparse prefill attention

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>

* rename

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>

---------

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>
2026-06-01 06:19:03 +08:00
cb7eb5c9b4 [MiniMax M3] Relocate sparse attention into the model dir + port decode top-k (#12)
Move the MiniMax M3 sparse-attention backend and Triton kernels out of the
shared attention tree into the model definition (mirroring deepseek_v4), split
the kernels by role, and tighten the backend, then port the dedicated split-K
decode top-k from the sglang reference.

Structure:
- vllm/v1/attention/backends/minimax_m3_sparse.py
    -> vllm/models/minimax_m3/common/sparse_attention.py
- vllm/v1/attention/ops/minimax_m3_sparse_ops.py split into
    common/ops/index_topk.py  (index-score + top-k kernels)
    common/ops/sparse_attn.py (block-sparse GQA attention kernels)
  Pure-Triton/cross-platform, so under common/ (not nvidia/).
- registry enum + doc generator RELEVANT_PATTERNS repointed; the auto-generated
  attention_backends.md now has a dedicated "MiniMax M3 Sparse Attention" section.

Backend cleanups:
- Drop redundant metadata fields (sparse-selection params, decode max_query_len,
  decode cu_seqlens_q/context_lens) and read them from the layer/impl instead.
- Impl ctor takes explicit named args instead of kwargs.get lookups.
- Cudagraph support UNIFORM_SINGLE_TOKEN_DECODE -> UNIFORM_BATCH; reorder
  threshold and decode buffer scale by 1 + num_speculative_tokens.
- Unify the per-request context-length buffer (max_num_batched_tokens).
- Declare bf16-only KV cache.

Decode top-k (ported from sglang minimax_sparse_ops/decode):
- Split-K top-k: per-chunk partial top-k (_topk_index_partial_kernel) + merge
  (_topk_index_merge_kernel), replacing the single-program prefill top-k reuse.
- init/local block forcing moved into the decode score kernel (matches sglang).
- Verified against a torch reference (selected-block sets match across varied
  seq lengths, with/without init+local forcing, and short sequences).

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-05-31 16:19:23 -04:00
Yongye Zhu 4f3c528941 [MiniMax M3] Enable decode CUDA graphs + dedicated split-K decode kernels (#7)
Two changes to the sparse-attention backend:

1. Full decode CUDA-graph support. The metadata builder now declares
   AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE and precomputes all
   per-step kernel args (cu_seqlens_q, prefix_lens, max_query_len,
   num_actual_tokens) in build(), removing the .item() host sync and the
   per-step torch.zeros/cumsum/diff from the impl forward. Derived decode
   prefix lengths are written into a persistent buffer so the captured
   graph reads stable addresses across replays.

2. Dedicated split-K decode kernels (mirroring the sglang reference)
   instead of reusing the prefill kernels with BLOCK_SIZE_Q=1, which left
   the GPU idle at decode (one query token per request). The index score
   now splits over seq blocks and the GQA attention splits over the
   selected top-k blocks with an LSE merge (flash-decoding). Chunk counts
   depend only on shape constants, so the grid is fixed within a CUDA
   graph.

Verified: a parity test against the prior (GSM8K 92.1) prefill-as-decode
path matches exactly on top-k selection and on attention output (bf16
noise) across seq lengths 128-2048.

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
2026-05-31 16:03:41 +00:00
Yongye ZhuandClaude Opus 4.8 0f854f78e8 [MiniMax M3] Implement sparse-attention forward + MXFP8 DeepGEMM MoE e2e
Port the MiniMax M3 lightning-indexer sparse attention to a Triton backend and
fix the MXFP8 DeepGEMM MoE path so the model runs end to end.

Attention (vllm/v1/attention/ops/minimax_m3_sparse_ops.py + backend):
- Triton kernels (paged, page == sparse block == 128): index block-score +
  bitonic top-k, and GQA block-sparse flash attention over the selected blocks.
- MiniMaxM3SparseImpl.forward: decode-first split, dispatching the same kernels
  per phase (a decode token is a 1-token prefill). Index and main caches use
  separate block tables.
- Dedicated MiniMaxM3IndexerBackend for the key-only index cache so the main
  GQA cache (num_kv_heads==1 at TP>=4) is not mistaken for the index layout.
- get_supported_kernel_block_sizes()==[128] (one sparse block per KV page).

MXFP8 DeepGEMM MoE:
- Prepare-phase activation quant emits float32 per-(1,32) group scales for the
  DeepGEMM backend (use_deep_gemm_packed_mxfp8 on the quant config), matching
  the FP8 128-block path with group=32.
- deepgemm_moe_permute / ep_scatter take a block_size so the activation-scale
  group (32) is honored through the expert permute.
- workspace_shapes uses the contiguous-layout M alignment (not block_shape[0],
  which is 1 for MXFP8 and under-sized the workspace).

GSM8K (5-shot, TP=4) flexible-extract 0.921 / strict 0.919.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Yongye Zhu <yongye@inferact.ai>

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
2026-05-31 16:03:39 +00:00
Yongye ZhuandClaude Opus 4.8 a3028cebbf [MiniMax M3] Add sparse-attention backend + merged attention layer
Scaffold the lightning-indexer sparse-attention path:

- MiniMaxM3SparseBackend (registered as MINIMAX_M3_SPARSE): block-sparse GQA
  backend; get_kv_cache_shape serves both the main K/V cache and the
  single-vector index-key side cache.
- MiniMaxM3IndexerCache: side KV cache for per-token index keys, key-only so it
  uses a single-vector MLAAttentionSpec rather than a K+V FullAttentionSpec.
- MiniMaxM3SparseMetadata (+ prefill/decode sub-metadata) and its builder,
  splitting the batch via split_decodes_and_prefills.
- MiniMaxM3SparseImpl: subclasses AttentionImplBase so it can take a custom
  forward(query, index_query, kv_cache, index_kv_cache); no alibi / sliding
  window / logits soft cap. forward is a stub pending the kernel port.

MiniMaxM3SparseAttention is merged into a single AttentionLayerBase: it owns the
projections, per-head QK norm and RoPE, binds the backend + impl, registers the
main K/V cache, and holds the index cache. Its forward computes q/k/v and the
index q/k, pre-inserts K/V and index-K into their caches, then calls the sparse
impl with only the queries.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Yongye Zhu <yongye@inferact.ai>

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
2026-05-31 16:02:56 +00:00
Yongye ZhuandClaude Opus 4.8 e327716282 [MiniMax M3] Drop dead index value/output projection branch
For M3, sparse_disable_index_value matches sparse_attention_freq exactly
([0,0,0,1,...,1]): the only layers with the flag unset (0-2) are the
non-sparse layers built as MiniMaxM3Attention. Every layer that constructs
MiniMaxM3SparseAttention therefore always disables the index value/output
projections, so index_{v,o}_proj are never created.

Remove the unreachable else branch, the disable_index_value parameter and
field, and the now-unused _disable_index_value_layer_ids helper.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Yongye Zhu <yongye@inferact.ai>

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
2026-05-31 16:02:10 +00:00
Yongye ZhuandClaude Opus 4.8 de1828b58f [MiniMax M3] Implement dense attention forward
Implement MiniMaxM3Attention.forward (dense path): qkv projection, split,
per-head QK norm (GemmaRMSNorm, qk_norm_type="per_head"), partial RoPE,
attention, and output projection. Mirrors the sglang reference dense path
and vLLM's canonical per-head-norm convention. attention_output_gate is
False for M3, so the gate branch is omitted.

The sparse attention forward (index branch) remains stubbed.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Yongye Zhu <yongye@inferact.ai>

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
2026-05-31 16:02:09 +00:00
Yongye ZhuandClaude Opus 4.8 6542ed479c [MiniMax M3] Add DeepGEMM MXFP8 MoE backend with swigluoai support
Extend DeepGemmExperts to support MXFP8 activations (FP8 e4m3 + UE8M0
1x32 block scales) via the grouped GEMM with recipe (1, 32), reusing the
oracle/fp8 weight-conversion path. Generalize
deepgemm_post_process_fp8_weight_block to derive the transform recipe
from the block shape ((1, 1, 32) for MXFP8) and accept uint8 E8M0 scales.

Unify the fused gated-activation+quant triton kernels around
y = (up + beta) * gate * sigmoid(alpha * gate): silu is alpha=1, beta=0
(bit-identical to before); swigluoai uses alpha/beta from config. Thread
gemm1_alpha/gemm1_beta from the FusedMoE layer through the MXFP8 quant
config into the kernels, and add swiglu_alpha/swiglu_beta to the layer
and MiniMax M3 config/model (beta sourced from config, not hardcoded).

Wire Fp8MoeBackend.DEEPGEMM into the MXFP8 oracle (selectable via
--moe-backend deep_gemm), resolving directly to DeepGemmExperts (the
Triton fallback cannot handle the 1x32 scheme). Advertise SWIGLUOAI in
_supports_activation so swigluoai selects DeepGEMM rather than falling
through to another backend; gate the MXFP8 scheme to Blackwell (SM100).

Verified on GB200: packed-kernel parity (silu defaults unchanged,
swigluoai matches torch ref), (1,32) weight-prep transform, and a TP=4
launch selecting the DEEPGEMM MXFP8 backend with full weight load (the
run then stops at the still-stubbed attention forward, as expected).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Yongye Zhu <yongye@inferact.ai>

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
2026-05-31 16:02:09 +00:00
Yongye ZhuandClaude Opus 4.8 485f33d796 [MiniMax M3] Port attention modules and wire up MXFP8 checkpoint loading
Port the weight-bearing dense (MiniMaxM3Attention) and sparse
(MiniMaxM3SparseAttention) attention modules so the checkpoint's
self_attn.* tensors map onto real params (forward still stubbed;
this targets weight loading). Add qkv stacked mapping and the
weight_scale_inv -> weight_scale remap in load_weights.

Load MiniMax-style MXFP8 checkpoints (quant_method: "mxfp8" +
ignored_layers) via the ModelOpt MXFP8 config: register "mxfp8" in
method_to_config and normalize the minimal checkpoint schema to the
ModelOpt schema in ModelOptMxFp8Config.from_config (same on-disk
format). Use setdefault for online shorthands so the checkpoint
config wins over the "mxfp8" online shorthand.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Yongye Zhu <yongye@inferact.ai>

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
2026-05-31 16:02:09 +00:00
Yongye ZhuandClaude Opus 4.8 1cad156ac5 [Model] MiniMax M3: implement MoE block + weight-name mapping
Implement the sigmoid-routed MoE block for MiniMax M3 and map module
names to the checkpoint so weight loading works for the ported modules.

MoE block (MiniMaxM3MoE):
- fp32 router via GateLinear (bf16 activations upcast to fp32; fp32
  weights and logits), matching minimax_m2/sglang precision.
- FusedTopKBiasRouter routing (scoring_func from config, sigmoid +
  e_score_correction_bias + renormalize), verified to match sglang's
  TopK (select-with-bias, weight-without-bias, routed_scaling on output).
- swigluoai activation (from config.hidden_act) + swiglu_limit; shared
  expert fused into FusedMoE so the shared partial is reduced with the
  routed output.

Weight loading:
- Name the MoE submodule `block_sparse_moe` (dense stays `mlp`) to match
  the checkpoint; decoder forward selects per layer.
- load_weights handles gate_up fusion (dense MLP + shared experts) and
  expert w1/w2/w3 -> w13/w2 fusion; wrappers delegate via
  AutoWeightsLoader, skipping vision/mm/mtp. Not-yet-ported modules
  (attention) are skipped until they land.

The expert GEMM/activation kernel correctness and attention/weight
loading for the remaining modules are not part of this change.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
2026-05-31 16:02:09 +00:00
Yongye ZhuandClaude Opus 4.8 99bd07204b [Model] Add MiniMax M3 text backbone skeleton + SwiGLU-OAI clamp activation
Port the MiniMax M3 (text backbone) into vLLM's custom model layout:

- Add MiniMaxM3SparseForCausalLM under vllm/models/minimax_m3/nvidia with the
  decoder/model/causal-LM wiring; attention and MoE bodies plus weight loading
  are left as stubs. Dense MiniMaxM3MLP is fully ported.
- Add MiniMaxM3SparseForConditionalGeneration as a minimal LM-routing wrapper
  (KimiK25-style init_vllm_registered_model on text_config) and register both
  architectures.
- Add MiniMaxM3Config (model_type minimax_m3_vl) wrapping MiniMaxM3TextConfig
  so config.get_text_config() extracts the backbone; register in the config
  registries.

Generalize silu_and_mul_with_clamp to SwiGLU-OAI:

- Add alpha (scales the activation's sigmoid) and beta (added to the
  non-activated half) to the CUDA kernel, ops.h, and torch_bindings schema.
  Defaults alpha=1.0, beta=0.0 are bitwise-identical to the previous
  silu(gate)*up, so existing callers (DeepSeek V4) are unaffected.
- SiluAndMulWithClamp(alpha, beta) used by MiniMaxM3MLP with alpha=swiglu_alpha,
  beta=1.0, matching the reference gate*sigmoid(alpha*gate)*(up+1).

AI assistance (Claude) was used for this change.

Signed-off-by: Yongye Zhu <yongye@inferact.ai>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
2026-05-31 16:02:09 +00:00
362 changed files with 26862 additions and 10206 deletions
+1 -3
View File
@@ -16,7 +16,6 @@ steps:
- tests/kernels/test_onednn.py
- tests/kernels/test_awq_int4_to_int8.py
- tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
- tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
commands:
- |
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
@@ -25,8 +24,7 @@ steps:
pytest -x -v -s tests/kernels/moe/test_cpu_quant_fused_moe.py
pytest -x -v -s tests/kernels/test_onednn.py
pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py
pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py"
pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py"
- label: CPU-Compatibility Tests
depends_on: []
@@ -37,8 +37,7 @@ function cpu_tests() {
pytest -x -v -s tests/kernels/test_onednn.py
pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
pytest -x -v -s tests/kernels/core/test_cpu_activation.py
pytest -x -v -s tests/kernels/moe/test_moe.py -k test_cpu_fused_moe_basic
pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py"
pytest -x -v -s tests/kernels/moe/test_moe.py -k test_cpu_fused_moe_basic"
# skip tests requiring model downloads if HF_TOKEN is not set
# due to rate-limits
@@ -1,39 +0,0 @@
#!/bin/bash
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
set -euo pipefail
REQUIREMENTS_FILE="${KV_CONNECTORS_REQUIREMENTS:-/vllm-workspace/requirements/kv_connectors.txt}"
uv pip install --system -r "${REQUIREMENTS_FILE}"
NIXL_METADATA=$(python3 - <<'PY'
import importlib.metadata as metadata
import torch
cuda_version = torch.version.cuda
if cuda_version is None:
raise SystemExit("torch.version.cuda is not set")
print(cuda_version.split(".", 1)[0], metadata.version("nixl"))
PY
)
read -r CUDA_MAJOR NIXL_VERSION <<<"${NIXL_METADATA}"
# nixl>=1.1.0 can install multiple CUDA wheel variants. Keep only the variant
# matching this CI image so nixl_ep_cpp links against the available libcudart.
uv pip uninstall --system nixl-cu12 nixl-cu13 2>/dev/null || true
uv pip install --system --no-deps "nixl-cu${CUDA_MAJOR}==${NIXL_VERSION}"
python3 - <<'PY'
import importlib.metadata as metadata
for package_name in ("nixl", "nixl-cu12", "nixl-cu13"):
try:
version = metadata.version(package_name)
except metadata.PackageNotFoundError:
version = "not installed"
print(f"{package_name}: {version}")
PY
+13 -15
View File
@@ -1238,11 +1238,14 @@ steps:
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/entrypoints/serve
- tests/entrypoints/rpc
- tests/entrypoints/serve/instrumentator
- tests/tool_use
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
- pytest -v -s entrypoints/serve/instrumentator
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
- pytest -v -s tool_use
- label: Entrypoints Integration (API Server openai - Part 1) # TBD
timeout_in_minutes: 180
@@ -1272,14 +1275,10 @@ steps:
- vllm/
- tests/entrypoints/openai
- tests/entrypoints/test_chat_utils
- tests/entrypoints/generate
- tests/tool_use
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
- pytest -v -s entrypoints/test_chat_utils.py
- pytest -v -s entrypoints/generate
- pytest -v -s tool_use
- label: Entrypoints Integration (API Server openai - Part 3) # TBD
timeout_in_minutes: 180
@@ -1369,7 +1368,7 @@ steps:
- vllm/platforms/rocm.py
commands:
- pytest -v -s entrypoints/openai/tool_parsers
- pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/offline_mode --ignore=entrypoints/openai --ignore=entrypoints/serve --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text --ignore=tests/entrypoints/generate
- pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/serve/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling
- label: OpenAI API correctness # TBD
timeout_in_minutes: 180
@@ -2746,11 +2745,14 @@ steps:
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/entrypoints/serve
- tests/entrypoints/rpc
- tests/entrypoints/serve/instrumentator
- tests/tool_use
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
- pytest -v -s entrypoints/serve/instrumentator
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
- pytest -v -s tool_use
- label: Entrypoints Integration (API Server openai - Part 1) # TBD
timeout_in_minutes: 180
@@ -2780,14 +2782,10 @@ steps:
- vllm/
- tests/entrypoints/openai
- tests/entrypoints/test_chat_utils
- tests/entrypoints/generate
- tests/tool_use
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
- pytest -v -s entrypoints/test_chat_utils.py
- pytest -v -s entrypoints/generate
- pytest -v -s tool_use
- label: Entrypoints Integration (API Server openai - Part 3) # TBD
timeout_in_minutes: 180
+9 -9
View File
@@ -11,7 +11,7 @@ steps:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: Distributed FlashInfer NixlConnector PD accuracy (4 GPUs)
key: distributed-flashinfer-nixlconnector-pd-accuracy-4-gpus
@@ -22,7 +22,7 @@ steps:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- FLASHINFER=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs)
@@ -34,7 +34,7 @@ steps:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs)
@@ -46,7 +46,7 @@ steps:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: Hybrid SSM NixlConnector PD accuracy tests (4 GPUs)
@@ -58,7 +58,7 @@ steps:
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
- tests/v1/kv_connector/nixl_integration/
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
- label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs)
@@ -73,7 +73,7 @@ steps:
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
- tests/v1/kv_connector/nixl_integration/
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
- label: NixlConnector PD + Spec Decode acceptance (2 GPUs)
@@ -87,7 +87,7 @@ steps:
- vllm/v1/worker/kv_connector_model_runner_mixin.py
- tests/v1/kv_connector/nixl_integration/
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash v1/kv_connector/nixl_integration/config_sweep_spec_decode_test.sh
- label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs)
@@ -102,5 +102,5 @@ steps:
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
- tests/v1/kv_connector/nixl_integration/
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
+8 -8
View File
@@ -11,7 +11,7 @@ steps:
- tests/entrypoints/
commands:
- pytest -v -s entrypoints/openai/tool_parsers
- pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/offline_mode --ignore=entrypoints/openai --ignore=entrypoints/serve --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text --ignore=tests/entrypoints/generate
- pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/serve/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text
- label: Entrypoints Integration (LLM)
key: entrypoints-integration-llm
@@ -60,13 +60,9 @@ steps:
- vllm/
- tests/entrypoints/openai
- tests/entrypoints/test_chat_utils
- tests/entrypoints/generate
- tests/tool_use
commands:
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
- pytest -v -s entrypoints/test_chat_utils.py
- pytest -v -s entrypoints/generate
- pytest -v -s tool_use
mirror:
amd:
device: mi325_1
@@ -102,11 +98,14 @@ steps:
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
- tests/entrypoints/serve
- tests/entrypoints/rpc
- tests/entrypoints/serve/instrumentator
- tests/tool_use
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
- pytest -v -s entrypoints/serve/instrumentator
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
- pytest -v -s tool_use
mirror:
amd:
device: mi325_1
@@ -154,5 +153,6 @@ steps:
source_file_dependencies:
- csrc/
- vllm/entrypoints/openai/
- vllm/model_executor/models/whisper.py
commands: # LMEval
- pytest -s entrypoints/openai/correctness/
+1 -1
View File
@@ -86,7 +86,7 @@ steps:
- tests/v1/metrics
- tests/entrypoints/openai/correctness/test_lmeval.py
commands:
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
# split the test to avoid interference
- pytest -v -s -m 'not cpu_test' v1/core
+4 -4
View File
@@ -45,19 +45,19 @@ steps:
- vllm/entrypoints/serve/
- vllm/v1/engine/
- tests/utils.py
# - tests/entrypoints/serve/dev/rpc/test_collective_rpc.py
# - tests/entrypoints/rpc/test_collective_rpc.py
- tests/entrypoints/serve/disagg/test_serving_tokens.py
- tests/entrypoints/serve/instrumentator/test_basic.py
- tests/entrypoints/serve/instrumentator/test_metrics.py
# - tests/entrypoints/serve/dev/test_sleep.py
# - tests/entrypoints/serve/instrumentator/test_sleep.py
commands:
- export VLLM_USE_RUST_FRONTEND=1
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
# - pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py
# - pytest -v -s entrypoints/rpc/test_collective_rpc.py
- pytest -v -s entrypoints/serve/instrumentator/test_basic.py -k "not show_version and not server_load"
- pytest -v -s entrypoints/serve/disagg/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow"
- pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist"
# - pytest -v -s entrypoints/serve/dev/test_sleep.py
# - pytest -v -s entrypoints/serve/instrumentator/test_sleep.py
- label: Rust Frontend Core Correctness
timeout_in_minutes: 30
+1 -1
View File
@@ -10,7 +10,7 @@ jobs:
runs-on: ubuntu-latest
steps:
- name: Add label
uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
with:
script: |
github.rest.issues.addLabels({
+3 -3
View File
@@ -14,7 +14,7 @@ jobs:
steps:
- name: Label issues based on keywords
id: label-step
uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
with:
script: |
// Configuration: Add new labels and keywords here
@@ -315,7 +315,7 @@ jobs:
- name: CC users for labeled issues
if: steps.label-step.outputs.labels_added != '[]'
uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
with:
script: |
// Configuration: Map labels to GitHub users to CC
@@ -392,7 +392,7 @@ jobs:
- name: Request missing ROCm info from issue author
if: contains(steps.label-step.outputs.labels_added, 'rocm') && contains(toJSON(github.event.issue.labels.*.name), 'bug')
uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
with:
script: |
const body = (context.payload.issue.body || '').toLowerCase();
+2 -2
View File
@@ -12,7 +12,7 @@ jobs:
runs-on: ubuntu-latest
steps:
- name: Update PR description
uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
with:
script: |
const { owner, repo } = context.repo;
@@ -55,7 +55,7 @@ jobs:
runs-on: ubuntu-latest
steps:
- name: Post welcome comment for first-time contributors
uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
with:
script: |
const { owner, repo } = context.repo;
+2 -2
View File
@@ -20,7 +20,7 @@ jobs:
runs-on: ubuntu-latest
steps:
- name: Check PR label and author merge count
uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0
uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8.0.0
with:
script: |
const { data: pr } = await github.rest.pulls.get({
@@ -49,7 +49,7 @@ jobs:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
- uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548 # v6.1.0
with:
python-version: "3.12"
- run: echo "::add-matcher::.github/workflows/matchers/actionlint.json"
+1 -1
View File
@@ -21,7 +21,7 @@ repos:
rev: v21.1.2
hooks:
- id: clang-format
exclude: 'csrc/(moe/topk_softmax_kernels.cu|libtorch_stable/quantization/gguf/(ggml-common.h|dequantize.cuh|vecdotq.cuh|mmq.cuh|mmvq.cuh))|vllm/third_party/.*'
exclude: 'csrc/(moe/topk_softmax_kernels.cu|quantization/gguf/(ggml-common.h|dequantize.cuh|vecdotq.cuh|mmq.cuh|mmvq.cuh))|vllm/third_party/.*'
types_or: [c++, cuda]
args: [--style=file, --verbose]
- repo: https://github.com/DavidAnson/markdownlint-cli2
View File
+3 -1
View File
@@ -315,7 +315,8 @@ set(VLLM_EXT_SRC
if(VLLM_GPU_LANG STREQUAL "CUDA")
list(APPEND VLLM_EXT_SRC
"csrc/minimax_reduce_rms_kernel.cu")
"csrc/minimax_reduce_rms_kernel.cu"
"csrc/minimax_m3_build_k2q_csr.cu")
SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
@@ -637,6 +638,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
"csrc/libtorch_stable/quantization/gguf/gguf_kernel.cu"
"csrc/libtorch_stable/pos_encoding_kernels.cu"
"csrc/libtorch_stable/fused_qknorm_rope_kernel.cu"
"csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu"
"csrc/libtorch_stable/layernorm_kernels.cu"
"csrc/libtorch_stable/layernorm_quant_kernels.cu"
"csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu"
+1 -31
View File
@@ -369,18 +369,6 @@ else()
add_compile_definitions(-DVLLM_NUMA_DISABLED)
endif()
# check if the pytorch wheel ships libopenblas.so.
set(VLLM_OPENBLAS_LIB "")
if (NOT ENABLE_X86_ISA)
file(GLOB _VLLM_TORCH_OPENBLAS_LIBS
"${TORCH_INSTALL_PREFIX}/lib/libopenblas*.so*")
# Note: we don't link openblas directly to _C extension, as it's available through libtorch.so
if (_VLLM_TORCH_OPENBLAS_LIBS)
list(GET _VLLM_TORCH_OPENBLAS_LIBS 0 VLLM_OPENBLAS_LIB)
message(STATUS "CPU OpenBLAS library: ${VLLM_OPENBLAS_LIB}")
endif()
endif()
#
# Generate CPU attention dispatch header
#
@@ -399,7 +387,6 @@ endif()
#
set(VLLM_EXT_SRC
"csrc/cpu/activation.cpp"
"csrc/cpu/sgl-kernels/fla.cpp"
"csrc/cpu/utils.cpp"
"csrc/cpu/spec_decode_utils.cpp"
"csrc/cpu/layernorm.cpp"
@@ -409,13 +396,6 @@ set(VLLM_EXT_SRC
"csrc/cpu/cpu_attn.cpp"
"csrc/cpu/torch_bindings.cpp")
if (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64" AND VLLM_RVV_VLEN AND
VLLM_RVV_VLEN GREATER 0 AND (RVV_FP16_FOUND OR RVV_BF16_FOUND))
set(VLLM_EXT_SRC
"csrc/cpu/cpu_wna16.cpp"
${VLLM_EXT_SRC})
endif()
if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
set(VLLM_EXT_SRC
"csrc/cpu/shm.cpp"
@@ -423,12 +403,6 @@ if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
${VLLM_EXT_SRC})
endif()
if (POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND)
set(VLLM_EXT_SRC
"csrc/cpu/shm.cpp"
${VLLM_EXT_SRC})
endif()
if(USE_ONEDNN)
set(VLLM_EXT_SRC
"csrc/cpu/dnnl_kernels.cpp"
@@ -437,6 +411,7 @@ endif()
if (ENABLE_X86_ISA)
set(VLLM_EXT_SRC_SGL
"csrc/cpu/sgl-kernels/fla.cpp"
"csrc/cpu/sgl-kernels/conv.cpp"
"csrc/cpu/sgl-kernels/gemm.cpp"
"csrc/cpu/sgl-kernels/gemm_int8.cpp"
@@ -448,7 +423,6 @@ if (ENABLE_X86_ISA)
"csrc/cpu/sgl-kernels/moe_fp8.cpp")
set(VLLM_EXT_SRC_AVX512
"csrc/cpu/sgl-kernels/fla.cpp"
"csrc/cpu/shm.cpp"
"csrc/cpu/cpu_wna16.cpp"
"csrc/cpu/cpu_fused_moe.cpp"
@@ -465,7 +439,6 @@ if (ENABLE_X86_ISA)
"csrc/moe/dynamic_4bit_int_moe_cpu.cpp")
set(VLLM_EXT_SRC_AVX2
"csrc/cpu/sgl-kernels/fla.cpp"
"csrc/cpu/utils.cpp"
"csrc/cpu/spec_decode_utils.cpp"
"csrc/cpu/cpu_attn.cpp"
@@ -539,9 +512,6 @@ else()
USE_SABI 3
WITH_SOABI
)
if (VLLM_OPENBLAS_LIB)
target_compile_definitions(_C PRIVATE VLLM_HAS_OPENBLAS)
endif()
endif()
message(STATUS "Enabling C extension.")
@@ -31,7 +31,7 @@ endif()
if(VLLM_FLASH_ATTN_SRC_DIR)
FetchContent_Declare(
vllm-flash-attn SOURCE_DIR
vllm-flash-attn SOURCE_DIR
${VLLM_FLASH_ATTN_SRC_DIR}
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
)
@@ -39,7 +39,7 @@ else()
FetchContent_Declare(
vllm-flash-attn
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
GIT_TAG dd62dac706b1cf7895bd99b18c6cb7e7e117ee25
GIT_TAG bce29425653ec0fbc579d329883030e832d15ada
GIT_PROGRESS TRUE
# Don't share the vllm-flash-attn build between build types
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
-51
View File
@@ -94,10 +94,6 @@ struct FP16Vec16 : public Vec<FP16Vec16> {
: reg(RVVI(__riscv_vle16_v_f16, LMUL_256)(
static_cast<const _Float16*>(ptr), VEC_ELEM_NUM)) {};
explicit FP16Vec16(const c10::Half v)
: reg(RVVI4(__riscv_vreinterpret_v_u16, LMUL_256, _f16, LMUL_256)(
RVVI(__riscv_vmv_v_x_u16, LMUL_256)(v.x, VEC_ELEM_NUM))) {};
explicit FP16Vec16(const FP32Vec16& vec);
void save(void* ptr) const {
@@ -169,9 +165,6 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
reinterpret_cast<const uint16_t*>(ptr), VEC_ELEM_NUM))) {};
explicit BF16Vec16(fixed_bf16x16_t data) : reg(data) {};
explicit BF16Vec16(const c10::BFloat16 v)
: reg(RVVI4(__riscv_vreinterpret_v_u16, LMUL_256, _bf16, LMUL_256)(
RVVI(__riscv_vmv_v_x_u16, LMUL_256)(v.x, VEC_ELEM_NUM))) {};
explicit BF16Vec16(const FP32Vec16&);
void save(void* ptr) const {
@@ -297,9 +290,6 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
}
reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_512)(tmp, 16);
}
explicit BF16Vec16(const c10::BFloat16 v)
: reg_fp32(RVVI(__riscv_vfmv_v_f_f32, LMUL_512)(static_cast<float>(v),
VEC_ELEM_NUM)) {}
explicit BF16Vec16(const FP32Vec16&);
void save(void* ptr) const {
float tmp[16];
@@ -639,19 +629,6 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
: reg(RVVI4(__riscv_vcreate_v_f32, LMUL_256, _f32, LMUL_512)(
data.reg, data.reg)) {};
explicit FP32Vec16(const FP32Vec16& data) : reg(data.reg) {};
explicit FP32Vec16(int64_t value, const FP32Vec16& lut) {
const uint64_t q_values = static_cast<uint64_t>(value);
auto packed = RVVI(__riscv_vmv_v_x_u64, LMUL_1024)(q_values, VEC_ELEM_NUM);
auto lane_ids = RVVI(__riscv_vid_v_u64, LMUL_1024)(VEC_ELEM_NUM);
auto shifts =
RVVI(__riscv_vsll_vx_u64, LMUL_1024)(lane_ids, 2, VEC_ELEM_NUM);
auto shifted =
RVVI(__riscv_vsrl_vv_u64, LMUL_1024)(packed, shifts, VEC_ELEM_NUM);
auto idx64 =
RVVI(__riscv_vand_vx_u64, LMUL_1024)(shifted, 0xF, VEC_ELEM_NUM);
auto idx32 = RVVI(__riscv_vnsrl_wx_u32, LMUL_512)(idx64, 0, VEC_ELEM_NUM);
reg = RVVI(__riscv_vrgather_vv_f32, LMUL_512)(lut.reg, idx32, VEC_ELEM_NUM);
}
explicit FP32Vec16(const FP16Vec16& v);
#ifdef __riscv_zvfbfmin
@@ -664,10 +641,6 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
explicit FP32Vec16(const BF16Vec16& v) : reg(v.reg_fp32) {};
#endif
// FP8 stub: dead code on RISC-V (fp8 KV cache is x86-only), needed for
// load_b_pair_vec template to compile on all platforms.
explicit FP32Vec16(const BF16Vec32&, int) : FP32Vec16() {}
FP32Vec16 operator+(const FP32Vec16& b) const {
return FP32Vec16(
RVVI(__riscv_vfadd_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
@@ -918,30 +891,6 @@ inline void fma(FP32Vec16& acc, const FP32Vec16& a, const FP32Vec16& b) {
acc = acc.fma(a, b);
}
template <typename VecT>
static void interleave_save_16b(const VecT& vec0, const VecT& vec1, void* ptr) {
alignas(64) uint16_t values0[VecT::VEC_ELEM_NUM];
alignas(64) uint16_t values1[VecT::VEC_ELEM_NUM];
vec0.save(values0);
vec1.save(values1);
auto* packed = reinterpret_cast<uint32_t*>(ptr);
for (int32_t i = 0; i < VecT::VEC_ELEM_NUM; ++i) {
packed[i] = static_cast<uint32_t>(values0[i]) |
(static_cast<uint32_t>(values1[i]) << 16);
}
}
static void interleave_save(const FP16Vec16& vec0, const FP16Vec16& vec1,
void* ptr) {
interleave_save_16b(vec0, vec1, ptr);
}
static void interleave_save(const BF16Vec16& vec0, const BF16Vec16& vec1,
void* ptr) {
interleave_save_16b(vec0, vec1, ptr);
}
#ifdef __riscv_zvfbfmin
template <>
inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
+1 -106
View File
@@ -89,35 +89,6 @@ struct BF16Vec8 : public Vec<BF16Vec8> {
}
};
struct FP16Vec16 : public Vec<FP16Vec16> {
constexpr static int VEC_ELEM_NUM = 16;
ss16x8x2_t reg;
explicit FP16Vec16(const void* ptr) {
reg.val[0] = (__vector signed short)vec_xl(0, (signed short*)ptr);
reg.val[1] = (__vector signed short)vec_xl(16, (signed short*)ptr);
}
explicit FP16Vec16(bool, const void* ptr) : FP16Vec16(ptr) {}
explicit FP16Vec16(const FP32Vec16&);
void save(void* ptr) const {
vec_xst(reg.val[0], 0, (signed short*)ptr);
vec_xst(reg.val[1], 16, (signed short*)ptr);
}
void save(void* ptr, int elem_num) const {
int num = std::max(0, std::min(elem_num, VEC_ELEM_NUM));
if (num <= 8) {
vec_xst_len(reg.val[0], (signed short*)ptr, num * 2);
} else {
vec_xst(reg.val[0], 0, (signed short*)ptr);
vec_xst_len(reg.val[1], (signed short*)ptr + 8, (num - 8) * 2);
}
}
};
struct BF16Vec16 : public Vec<BF16Vec16> {
constexpr static int VEC_ELEM_NUM = 16;
@@ -129,8 +100,6 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
reg.val[1] = (__vector signed short)vec_xl(16, (signed short*)ptr);
}
explicit BF16Vec16(bool, const void* ptr) : BF16Vec16(ptr) {}
explicit BF16Vec16(const FP32Vec16&);
void save(void* ptr) const {
@@ -410,8 +379,6 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
reg.val[3] = vec_xl(48, ptr);
}
explicit FP32Vec16(bool, const float* ptr) : FP32Vec16(ptr) {}
explicit FP32Vec16(f32x4x4_t data) : reg(data) {}
explicit FP32Vec16(const FP32Vec16& data) {
@@ -435,7 +402,6 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
reg.val[3] = data.reg.val[1];
}
explicit FP32Vec16(const FP16Vec16& v);
explicit FP32Vec16(const BF16Vec16& v) {
reg.val[0] = (__vector float)vec_mergeh(zero, v.reg.val[0]);
reg.val[1] = (__vector float)vec_mergel(zero, v.reg.val[0]);
@@ -769,40 +735,6 @@ inline BF16Vec8::BF16Vec8(const FP32Vec8& v) {
#endif
}
inline FP16Vec16::FP16Vec16(const FP32Vec16& v) {
alignas(16) float temp_fp32[16];
alignas(16) c10::Half temp_fp16[16];
vec_xst(v.reg.val[0], 0, temp_fp32);
vec_xst(v.reg.val[1], 16, temp_fp32);
vec_xst(v.reg.val[2], 32, temp_fp32);
vec_xst(v.reg.val[3], 48, temp_fp32);
for (int i = 0; i < 16; i++) {
temp_fp16[i] = c10::Half(temp_fp32[i]);
}
reg.val[0] = (__vector signed short)vec_xl(0, (signed short*)temp_fp16);
reg.val[1] = (__vector signed short)vec_xl(16, (signed short*)temp_fp16);
}
inline FP32Vec16::FP32Vec16(const FP16Vec16& v) {
alignas(16) c10::Half temp_fp16[16];
alignas(16) float temp_fp32[16];
vec_xst(v.reg.val[0], 0, (signed short*)temp_fp16);
vec_xst(v.reg.val[1], 16, (signed short*)temp_fp16);
for (int i = 0; i < 16; i++) {
temp_fp32[i] = float(temp_fp16[i]);
}
reg.val[0] = vec_xl(0, temp_fp32);
reg.val[1] = vec_xl(16, temp_fp32);
reg.val[2] = vec_xl(32, temp_fp32);
reg.val[3] = vec_xl(48, temp_fp32);
}
inline BF16Vec16::BF16Vec16(const FP32Vec16& v) {
#ifdef _ARCH_PWR10
__vector signed short ret[4];
@@ -862,43 +794,6 @@ inline void prefetch(const void* addr) {
__asm__ __volatile__("dcbt 0, %0" : : "r"(addr) : "memory");
}
struct INT8Vec64 {
__vector signed char data[4];
INT8Vec64() = default;
explicit INT8Vec64(const int8_t* ptr) {
data[0] = vec_xl(0, ptr);
data[1] = vec_xl(16, ptr);
data[2] = vec_xl(32, ptr);
data[3] = vec_xl(48, ptr);
}
explicit INT8Vec64(bool, const int8_t* ptr) : INT8Vec64(ptr) {}
void save(int8_t* ptr) const {
vec_xst(data[0], 0, ptr);
vec_xst(data[1], 16, ptr);
vec_xst(data[2], 32, ptr);
vec_xst(data[3], 48, ptr);
}
void save(int8_t* ptr, int elem_num) const {
if (elem_num <= 0) return;
int full_vecs = elem_num / 16;
for (int i = 0; i < full_vecs && i < 4; i++) {
vec_xst(data[i], i * 16, ptr);
}
int remaining = elem_num % 16;
if (remaining > 0 && full_vecs < 4) {
vec_xst_len(data[full_vecs], ptr + full_vecs * 16, remaining);
}
}
void nt_save(int8_t* ptr) const { save(ptr); }
};
} // namespace vec_op
}; // namespace vec_op
#endif
-82
View File
@@ -1,82 +0,0 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
#include <ATen/native/CPUBlas.h>
// Unlike brgemm, PyTorch does not publicly expose at::native::cpublas::gemm
// If OpenBLS is available in the PyTorch wheel, we rely on it for fast
// bf16:bf16->fp32 GEMMs Otherwise, we fall back to PyTorch reference BLAS path.
#if defined(VLLM_HAS_OPENBLAS)
extern "C" void sbgemm_(char* transa, char* transb, int* m, int* n, int* k,
float* alpha, const at::BFloat16* a, int* lda,
const at::BFloat16* b, int* ldb, float* beta, float* c,
int* ldc);
extern "C" void sgemm_(char* transa, char* transb, int* m, int* n, int* k,
float* alpha, const float* a, int* lda, const float* b,
int* ldb, float* beta, float* c, int* ldc);
inline char blas_transpose(at::native::TransposeType trans) {
switch (trans) {
case at::native::TransposeType::NoTranspose:
return 'n';
case at::native::TransposeType::Transpose:
return 't';
case at::native::TransposeType::ConjTranspose:
return 'c';
}
return 'n';
}
inline void blas_gemm(at::native::TransposeType transa,
at::native::TransposeType transb, int64_t m, int64_t n,
int64_t k, float alpha, const at::BFloat16* a,
int64_t lda, const at::BFloat16* b, int64_t ldb,
float beta, float* c, int64_t ldc) {
char transa_ = blas_transpose(transa);
char transb_ = blas_transpose(transb);
int m_ = static_cast<int>(m);
int n_ = static_cast<int>(n);
int k_ = static_cast<int>(k);
int lda_ = static_cast<int>(lda);
int ldb_ = static_cast<int>(ldb);
int ldc_ = static_cast<int>(ldc);
sbgemm_(&transa_, &transb_, &m_, &n_, &k_, &alpha, a, &lda_, b, &ldb_, &beta,
c, &ldc_);
}
inline void blas_gemm(at::native::TransposeType transa,
at::native::TransposeType transb, int64_t m, int64_t n,
int64_t k, float alpha, const float* a, int64_t lda,
const float* b, int64_t ldb, float beta, float* c,
int64_t ldc) {
char transa_ = blas_transpose(transa);
char transb_ = blas_transpose(transb);
int m_ = static_cast<int>(m);
int n_ = static_cast<int>(n);
int k_ = static_cast<int>(k);
int lda_ = static_cast<int>(lda);
int ldb_ = static_cast<int>(ldb);
int ldc_ = static_cast<int>(ldc);
sgemm_(&transa_, &transb_, &m_, &n_, &k_, &alpha, a, &lda_, b, &ldb_, &beta,
c, &ldc_);
}
inline void blas_gemm(at::native::TransposeType, at::native::TransposeType,
int64_t, int64_t, int64_t, float, const at::Half*,
int64_t, const at::Half*, int64_t, float, float*,
int64_t) {
TORCH_CHECK(false, "CPU OpenBLAS hgemm is not available.");
}
#else
template <typename scalar_t>
inline void blas_gemm(at::native::TransposeType transa,
at::native::TransposeType transb, int64_t m, int64_t n,
int64_t k, float alpha, const scalar_t* a, int64_t lda,
const scalar_t* b, int64_t ldb, float beta, float* c,
int64_t ldc) {
auto gemm = at::native::cpublas::gemm_no_downcast_stub.DEFAULT;
gemm(c10::CppTypeToScalarType<scalar_t>::value, transa, transb, m, n, k,
at::Scalar(alpha), a, lda, b, ldb, at::Scalar(beta), c, ldc);
}
#endif
+141 -278
View File
@@ -301,42 +301,25 @@ void chunk_gated_delta_rule_kernel_impl(
// attn = k_beta @ key.transpose(-1, -2)
// attn: [B, HV, num_chunk, chunk_size, chunk_size]
// transpose and pack for key
if constexpr (brgemm_supported()) {
pack_vnni<scalar_t>(
/* dst */ k_transpose,
/* src */ curr_k_pad,
/* N */ chunk_size,
/* K */ qk_head_size,
/* ld_src */ qk_head_size,
/* ld_dst */ chunk_size);
// k_beta @ key.transpose(-1, -2)
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ chunk_size,
/* K */ qk_head_size,
/* lda */ qk_head_size,
/* ldb */ chunk_size,
/* ldc */ chunk_size,
/* add_C */ false,
/* A */ curr_k_beta,
/* B */ k_transpose,
/* C */ curr_attn);
} else {
blas_gemm(
at::native::TransposeType::Transpose,
at::native::TransposeType::NoTranspose,
chunk_size,
chunk_size,
qk_head_size,
1.0f,
curr_k_pad,
qk_head_size,
curr_k_beta,
qk_head_size,
0.0f,
curr_attn,
chunk_size);
}
pack_vnni<scalar_t>(
/* dst */ k_transpose,
/* src */ curr_k_pad,
/* N */ chunk_size,
/* K */ qk_head_size,
/* ld_src */ qk_head_size,
/* ld_dst */ chunk_size);
// k_beta @ key.transpose(-1, -2)
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ chunk_size,
/* K */ qk_head_size,
/* lda */ qk_head_size,
/* ldb */ chunk_size,
/* ldc */ chunk_size,
/* add_C */ false,
/* A */ curr_k_beta,
/* B */ k_transpose,
/* C */ curr_attn);
// attn = attn * decay_mask
for (int64_t m = 0; m < chunk_size; m++) {
at::vec::map2<float>(
@@ -430,42 +413,25 @@ void chunk_gated_delta_rule_kernel_impl(
// k_beta_g = k_beta * g: [B, HV, num_chunk, chunk_size, EK]
// k_cumdecay: [B, HV, num_chunk, chunk_size, EK]
// pack for value
if constexpr (brgemm_supported()) {
pack_vnni2<scalar_t>(
/* dst */ v_pack,
/* src */ curr_v_beta,
/* N */ chunk_size,
/* K */ v_head_size,
/* ld_src */ v_head_size,
/* ld_dst */ v_head_size);
// value = attn @ v_beta
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ v_head_size,
/* K */ chunk_size,
/* lda */ chunk_size,
/* ldb */ v_head_size,
/* ldc */ v_head_size,
/* add_C */ false,
/* A */ curr_attn_reduced,
/* B */ v_pack,
/* C */ curr_value);
} else {
blas_gemm(
at::native::TransposeType::NoTranspose,
at::native::TransposeType::NoTranspose,
v_head_size,
chunk_size,
chunk_size,
1.0f,
curr_v_beta,
v_head_size,
curr_attn_reduced,
chunk_size,
0.0f,
curr_value,
v_head_size);
}
pack_vnni2<scalar_t>(
/* dst */ v_pack,
/* src */ curr_v_beta,
/* N */ chunk_size,
/* K */ v_head_size,
/* ld_src */ v_head_size,
/* ld_dst */ v_head_size);
// value = attn @ v_beta
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ v_head_size,
/* K */ chunk_size,
/* lda */ chunk_size,
/* ldb */ v_head_size,
/* ldc */ v_head_size,
/* add_C */ false,
/* A */ curr_attn_reduced,
/* B */ v_pack,
/* C */ curr_value);
// k_beta_g = k_beta * g.exp().unsqueeze(-1)
for (int64_t j = 0; j < chunk_size; j++) {
int64_t i = 0;
@@ -479,42 +445,25 @@ void chunk_gated_delta_rule_kernel_impl(
}
}
// pack for k_beta_g
if constexpr (brgemm_supported()) {
pack_vnni2<scalar_t>(
/* dst */ k_beta_g_pack,
/* src */ k_beta_g,
/* N */ chunk_size,
/* K */ qk_head_size,
/* ld_src */ qk_head_size,
/* ld_dst */ qk_head_size);
// k_cumdecay = attn @ k_beta_g
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ qk_head_size,
/* K */ chunk_size,
/* lda */ chunk_size,
/* ldb */ qk_head_size,
/* ldc */ qk_head_size,
/* add_C */ false,
/* A */ curr_attn_reduced,
/* B */ k_beta_g_pack,
/* C */ k_cumdecay);
} else {
blas_gemm(
at::native::TransposeType::NoTranspose,
at::native::TransposeType::NoTranspose,
qk_head_size,
chunk_size,
chunk_size,
1.0f,
k_beta_g,
qk_head_size,
curr_attn_reduced,
chunk_size,
0.0f,
k_cumdecay,
qk_head_size);
}
pack_vnni2<scalar_t>(
/* dst */ k_beta_g_pack,
/* src */ k_beta_g,
/* N */ chunk_size,
/* K */ qk_head_size,
/* ld_src */ qk_head_size,
/* ld_dst */ qk_head_size);
// k_cumdecay = attn @ k_beta_g
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ qk_head_size,
/* K */ chunk_size,
/* lda */ chunk_size,
/* ldb */ qk_head_size,
/* ldc */ qk_head_size,
/* add_C */ false,
/* A */ curr_attn_reduced,
/* B */ k_beta_g_pack,
/* C */ k_cumdecay);
for (int i = 0; i < chunk_size; i++) {
at::vec::map<scalar_t>(
[](fVec x) { return x; },
@@ -602,42 +551,25 @@ void chunk_gated_delta_rule_kernel_impl(
// attn_i = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask, 0)
// k_transpose_i = k_i.transpose(-1, -2)
if constexpr (brgemm_supported()) {
pack_vnni<scalar_t>(
/* dst */ k_transpose_i,
/* src */ k_i,
/* N */ chunk_size,
/* K */ qk_head_size,
/* ld_src */ qk_head_size,
/* ld_dst */ chunk_size);
// attn_i = q_i @ k_transpose_i
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ chunk_size,
/* K */ qk_head_size,
/* lda */ qk_head_size,
/* ldb */ chunk_size,
/* ldc */ chunk_size,
/* add_C */ false,
/* A */ q_i,
/* B */ k_transpose_i,
/* C */ attn_i);
} else {
blas_gemm(
at::native::TransposeType::Transpose,
at::native::TransposeType::NoTranspose,
chunk_size,
chunk_size,
qk_head_size,
1.0f,
k_i,
qk_head_size,
q_i,
qk_head_size,
0.0f,
attn_i,
chunk_size);
}
pack_vnni<scalar_t>(
/* dst */ k_transpose_i,
/* src */ k_i,
/* N */ chunk_size,
/* K */ qk_head_size,
/* ld_src */ qk_head_size,
/* ld_dst */ chunk_size);
// attn_i = q_i @ k_transpose_i
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ chunk_size,
/* K */ qk_head_size,
/* lda */ qk_head_size,
/* ldb */ chunk_size,
/* ldc */ chunk_size,
/* add_C */ false,
/* A */ q_i,
/* B */ k_transpose_i,
/* C */ attn_i);
// attn_i = attn_i * decay_mask_i
for (int64_t m = 0; m < chunk_size; m++) {
auto attn_i_m = attn_i + m * chunk_size;
@@ -677,45 +609,28 @@ void chunk_gated_delta_rule_kernel_impl(
}
// pack for curr_last_recurrent_state
if constexpr (brgemm_supported()) {
pack_vnni2<scalar_t>(
/* dst */ curr_last_recurrent_state_pack_reduced,
/* src */ curr_last_recurrent_state_reduced,
/* N */ qk_head_size,
/* K */ v_head_size,
/* ld_src */ v_head_size,
/* ld_dst */ v_head_size);
pack_vnni2<scalar_t>(
/* dst */ curr_last_recurrent_state_pack_reduced,
/* src */ curr_last_recurrent_state_reduced,
/* N */ qk_head_size,
/* K */ v_head_size,
/* ld_src */ v_head_size,
/* ld_dst */ v_head_size);
// v_prime = k_cumdecay_i @ curr_last_recurrent_state: [chunk_size, EV]
// k_cumdecay_i: [chunk_size, EK]
// curr_last_recurrent_state: [EK, EV]
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ v_head_size,
/* K */ qk_head_size,
/* lda */ qk_head_size,
/* ldb */ v_head_size,
/* ldc */ v_head_size,
/* add_C */ false,
/* A */ k_cumdecay_i_reduced,
/* B */ curr_last_recurrent_state_pack_reduced,
/* C */ v_prime);
} else {
blas_gemm(
at::native::TransposeType::NoTranspose,
at::native::TransposeType::NoTranspose,
v_head_size,
chunk_size,
qk_head_size,
1.0f,
curr_last_recurrent_state_reduced,
v_head_size,
k_cumdecay_i_reduced,
qk_head_size,
0.0f,
v_prime,
v_head_size);
}
// v_prime = k_cumdecay_i @ curr_last_recurrent_state: [chunk_size, EV]
// k_cumdecay_i: [chunk_size, EK]
// curr_last_recurrent_state: [EK, EV]
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ v_head_size,
/* K */ qk_head_size,
/* lda */ qk_head_size,
/* ldb */ v_head_size,
/* ldc */ v_head_size,
/* add_C */ false,
/* A */ k_cumdecay_i_reduced,
/* B */ curr_last_recurrent_state_pack_reduced,
/* C */ v_prime);
// v_new = v_prime = v_i - v_prime
// v_i: [chunk_size, EV]
@@ -748,75 +663,41 @@ void chunk_gated_delta_rule_kernel_impl(
}
// attn_inter = qg @ curr_last_recurrent_state: [chunk_size, EV]
// curr_last_recurrent_state: [EK, EV]
if constexpr (brgemm_supported()) {
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ v_head_size,
/* K */ qk_head_size,
/* lda */ qk_head_size,
/* ldb */ v_head_size,
/* ldc */ v_head_size,
/* add_C */ false,
/* A */ qg,
/* B */ curr_last_recurrent_state_pack_reduced,
/* C */ attn_inter);
} else {
blas_gemm(
at::native::TransposeType::NoTranspose,
at::native::TransposeType::NoTranspose,
v_head_size,
chunk_size,
qk_head_size,
1.0f,
curr_last_recurrent_state_reduced,
v_head_size,
qg,
qk_head_size,
0.0f,
attn_inter,
v_head_size);
}
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ v_head_size,
/* K */ qk_head_size,
/* lda */ qk_head_size,
/* ldb */ v_head_size,
/* ldc */ v_head_size,
/* add_C */ false,
/* A */ qg,
/* B */ curr_last_recurrent_state_pack_reduced,
/* C */ attn_inter);
// core_attn_out[:, :, i] = attn_inter + attn_i @ v_new
// pack for v_prime
if constexpr (brgemm_supported()) {
pack_vnni2<scalar_t>(
/* dst */ v_prime_pack_reduced,
/* src */ v_prime_reduced,
/* N */ chunk_size,
/* K */ v_head_size,
/* ld_src */ v_head_size,
/* ld_dst */ v_head_size);
// attn_inter = attn_inter + attn_i @ v_new: [chunk_size, EV]
// attn_i: [chunk_size, chunk_size]
// v_new: [chunk_size, EV]
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ v_head_size,
/* K */ chunk_size,
/* lda */ chunk_size,
/* ldb */ v_head_size,
/* ldc */ v_head_size,
/* add_C */ true,
/* A */ attn_i_reduced,
/* B */ v_prime_pack_reduced,
/* C */ attn_inter);
} else {
blas_gemm(
at::native::TransposeType::NoTranspose,
at::native::TransposeType::NoTranspose,
v_head_size,
chunk_size,
chunk_size,
1.0f,
v_prime_reduced,
v_head_size,
attn_i_reduced,
chunk_size,
1.0f,
attn_inter,
v_head_size);
}
pack_vnni2<scalar_t>(
/* dst */ v_prime_pack_reduced,
/* src */ v_prime_reduced,
/* N */ chunk_size,
/* K */ v_head_size,
/* ld_src */ v_head_size,
/* ld_dst */ v_head_size);
// attn_inter = attn_inter + attn_i @ v_new: [chunk_size, EV]
// attn_i: [chunk_size, chunk_size]
// v_new: [chunk_size, EV]
at::native::cpublas::brgemm(
/* M */ chunk_size,
/* N */ v_head_size,
/* K */ chunk_size,
/* lda */ chunk_size,
/* ldb */ v_head_size,
/* ldc */ v_head_size,
/* add_C */ true,
/* A */ attn_i_reduced,
/* B */ v_prime_pack_reduced,
/* C */ attn_inter);
// core_attn_out[:, :, i] = attn_inter
for (int64_t m = 0; m < chunk_size; m++) {
@@ -881,34 +762,17 @@ void chunk_gated_delta_rule_kernel_impl(
/* ld_dst */ chunk_size);
// kgv = kg.transpose(-1, -2) @ v_new
// v_new: [chunk_size, EV]
if constexpr (brgemm_supported()) {
at::native::cpublas::brgemm(
/* M */ qk_head_size,
/* N */ v_head_size,
/* K */ chunk_size,
/* lda */ chunk_size,
/* ldb */ v_head_size,
/* ldc */ v_head_size,
/* add_C */ false,
/* A */ kg_transpose,
/* B */ v_prime_pack_reduced,
/* C */ kgv);
} else {
blas_gemm(
at::native::TransposeType::NoTranspose,
at::native::TransposeType::NoTranspose,
v_head_size,
qk_head_size,
chunk_size,
1.0f,
v_prime_reduced,
v_head_size,
kg_transpose,
chunk_size,
0.0f,
kgv,
v_head_size);
}
at::native::cpublas::brgemm(
/* M */ qk_head_size,
/* N */ v_head_size,
/* K */ chunk_size,
/* lda */ chunk_size,
/* ldb */ v_head_size,
/* ldc */ v_head_size,
/* add_C */ false,
/* A */ kg_transpose,
/* B */ v_prime_pack_reduced,
/* C */ kgv);
// last_recurrent_state = 1) + 2)
for (int64_t m = 0; m < qk_head_size; m++) {
at::vec::map2<float>(
@@ -1057,8 +921,7 @@ void fused_sigmoid_gating_delta_rule_update_kernel_impl(
float k_scale = use_qk_l2norm_in_kernel ? qk_scale_buf[k_scale_offset] : 1.0f;
int64_t v_offset = si * v_strideS + bi * v_strideB + ni * v_strideH;
int64_t o_offset = ((bi * seq_len + si) * v_num_heads + ni) * v_head_dim;
// See: https://github.com/sgl-project/sglang/pull/26634
float beta_val = 1 / (1 + std::exp(-b_ptr[bi * v_num_heads + ni]));
float beta_val = 1 / (1 + std::exp(-b_ptr[ni]));
fVec beta_vec = fVec(beta_val);
int64_t dvi = 0;
for (; dvi <= v_head_dim - VecSize; dvi += VecSize) {
+7 -18
View File
@@ -4,12 +4,9 @@
// clang-format off
#pragma once
#include "common.h"
#include "blas_gemm.h"
#include <ATen/native/CPUBlas.h>
#if defined(__AVX512F__) && defined(__AVX512BF16__) && defined(__AMX_BF16__)
#define CPU_CAPABILITY_AVX512
#endif
#include "common.h"
// amx-bf16
#define TILE_M 16
@@ -24,39 +21,31 @@ constexpr int block_size_n() {
return 2 * TILE_N;
}
constexpr bool brgemm_supported() {
#if defined(CPU_CAPABILITY_AVX512)
return true;
#else
return false;
#endif
}
// define threshold using brgemm (intel AMX)
template <typename T>
inline bool can_use_brgemm(int M);
template <>
inline bool can_use_brgemm<at::BFloat16>(int M) {
return brgemm_supported() && M > 4;
return M > 4;
}
template <>
inline bool can_use_brgemm<at::Half>(int M) {
return brgemm_supported();
return true;
}
// this requires PyTorch 2.7 or above
template <>
inline bool can_use_brgemm<int8_t>(int M) {
return brgemm_supported() && M > 4;
return M > 4;
}
template <>
inline bool can_use_brgemm<uint8_t>(int M) {
return brgemm_supported() && M > 4;
return M > 4;
}
template <>
inline bool can_use_brgemm<at::Float8_e4m3fn>(int M) {
return brgemm_supported() && M > 4;
return M > 4;
}
// work around compiler internal error
-2
View File
@@ -11,9 +11,7 @@
#include <ATen/cpu/vec/functional.h>
#include <ATen/cpu/vec/vec.h>
#if defined(CPU_CAPABILITY_AVX512)
#include <immintrin.h>
#endif
namespace {
using namespace at::vec;
+8 -10
View File
@@ -5,7 +5,7 @@
#include <sys/stat.h>
#include <unistd.h>
#if defined(__aarch64__) || defined(__powerpc64__)
#ifdef __aarch64__
#include <atomic>
#endif
@@ -38,7 +38,7 @@ struct KernelVecType<c10::Half> {
};
struct ThreadSHMContext {
#if defined(__aarch64__) || defined(__powerpc64__)
#ifdef __aarch64__
// memory model is weaker on AArch64, so we use atomic variables for
// consumer (load-acquire) and producer (store-release) to make sure
// that a stamp cannot be ready before the corresponding data is ready.
@@ -75,7 +75,7 @@ struct ThreadSHMContext {
TORCH_CHECK(group_size <= MAX_SHM_RANK_NUM);
TORCH_CHECK((size_t)this % 64 == 0);
TORCH_CHECK((size_t)thread_shm_ptr % 64 == 0);
#if defined(__aarch64__) || defined(__powerpc64__)
#ifdef __aarch64__
_curr_thread_stamp[0].store(1, std::memory_order_relaxed);
_curr_thread_stamp[1].store(1, std::memory_order_relaxed);
_ready_thread_stamp[0].store(0, std::memory_order_relaxed);
@@ -124,7 +124,7 @@ struct ThreadSHMContext {
}
char get_curr_stamp(int idx) const {
#if defined(__aarch64__) || defined(__powerpc64__)
#ifdef __aarch64__
return _curr_thread_stamp[idx].load(std::memory_order_acquire);
#else
return _curr_thread_stamp[idx];
@@ -132,7 +132,7 @@ struct ThreadSHMContext {
}
char get_ready_stamp(int idx) const {
#if defined(__aarch64__) || defined(__powerpc64__)
#ifdef __aarch64__
return _ready_thread_stamp[idx].load(std::memory_order_acquire);
#else
return _ready_thread_stamp[idx];
@@ -140,7 +140,7 @@ struct ThreadSHMContext {
}
void next_stamp() {
#if defined(__aarch64__) || defined(__powerpc64__)
#ifdef __aarch64__
_curr_thread_stamp[local_stamp_buffer_idx].fetch_add(
1, std::memory_order_release);
#else
@@ -150,7 +150,7 @@ struct ThreadSHMContext {
}
void commit_ready_stamp() {
#if defined(__aarch64__) || defined(__powerpc64__)
#ifdef __aarch64__
_ready_thread_stamp[local_stamp_buffer_idx].store(
_curr_thread_stamp[local_stamp_buffer_idx].load(
std::memory_order_relaxed),
@@ -186,10 +186,8 @@ struct ThreadSHMContext {
break;
}
++_spinning_count;
#if defined(__aarch64__)
#ifdef __aarch64__
__asm__ __volatile__("yield");
#elif defined(__powerpc64__)
__asm__ __volatile__("or 1,1,1");
#else
_mm_pause();
#endif // __aarch64__
+21 -22
View File
@@ -378,8 +378,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
#endif
// SHM CCL
#if defined(__AVX512F__) || (defined(__aarch64__) && !defined(__APPLE__)) || \
defined(__powerpc64__)
#if defined(__AVX512F__) || (defined(__aarch64__) && !defined(__APPLE__))
ops.def(
"init_shm_manager(str name, int group_size, int rank, int thread_num) -> "
"int",
@@ -448,25 +447,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
"bool is_vnni) -> Tensor");
ops.impl("fp8_scaled_mm_cpu", torch::kCPU, &fp8_scaled_mm_cpu);
// Adapted from sglang: casual_conv1d kernels
ops.def("causal_conv1d_weight_pack(Tensor weight) -> Tensor");
ops.impl("causal_conv1d_weight_pack", torch::kCPU,
&causal_conv1d_weight_pack);
ops.def(
"causal_conv1d_fwd_cpu(Tensor x, Tensor weight, Tensor? bias, Tensor? "
"conv_states, Tensor? query_start_loc,"
"Tensor? cache_indices, Tensor? has_initial_state, bool silu_activation, "
"int pad_slot_id, bool is_vnni) -> "
"Tensor");
ops.impl("causal_conv1d_fwd_cpu", torch::kCPU, &causal_conv1d_fwd_cpu);
ops.def(
"causal_conv1d_update_cpu(Tensor x, Tensor(a!) conv_states, Tensor "
"weight, Tensor? bias, bool silu_activation,"
"Tensor? cache_seqlens, Tensor? conv_state_indices, int pad_slot_id, "
"bool is_vnni) -> Tensor");
ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu);
#endif
// Adapted from sglang: GDN kernels
ops.def(
"chunk_gated_delta_rule_cpu(Tensor query, Tensor key, Tensor value, "
@@ -490,6 +470,25 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
"-> (Tensor, Tensor)");
ops.impl("fused_gdn_gating_cpu", torch::kCPU, &fused_gdn_gating_cpu);
// Adapted from sglang: casual_conv1d kernels
ops.def("causal_conv1d_weight_pack(Tensor weight) -> Tensor");
ops.impl("causal_conv1d_weight_pack", torch::kCPU,
&causal_conv1d_weight_pack);
ops.def(
"causal_conv1d_fwd_cpu(Tensor x, Tensor weight, Tensor? bias, Tensor? "
"conv_states, Tensor? query_start_loc,"
"Tensor? cache_indices, Tensor? has_initial_state, bool silu_activation, "
"int pad_slot_id, bool is_vnni) -> "
"Tensor");
ops.impl("causal_conv1d_fwd_cpu", torch::kCPU, &causal_conv1d_fwd_cpu);
ops.def(
"causal_conv1d_update_cpu(Tensor x, Tensor(a!) conv_states, Tensor "
"weight, Tensor? bias, bool silu_activation,"
"Tensor? cache_seqlens, Tensor? conv_state_indices, int pad_slot_id, "
"bool is_vnni) -> Tensor");
ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu);
#endif
// CPU attention kernels
ops.def(
"get_scheduler_metadata(int num_req, int num_heads_q, int num_heads_kv, "
@@ -519,7 +518,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
ops.def("dynamic_per_token_scaled_fp8_quant() -> ()", placeholder_op);
// WNA16
#if defined(__AVX512F__) || defined(__riscv_v)
#if defined(__AVX512F__)
ops.def(
"cpu_gemm_wna16(Tensor input, Tensor q_weight, Tensor(a2!) output, "
"Tensor scales, Tensor? zeros, Tensor? g_idx, Tensor? bias, SymInt "
+75 -41
View File
@@ -10,11 +10,20 @@
namespace vllm {
template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
// `alpha` and `beta` are applied to opposite operands:
// - alpha lives INSIDE the activation (the activated half): the gated
// activation computes act_half * sigmoid(alpha * act_half).
// - beta is added to the OTHER (non-activated) half before the multiply.
// So the result is always ACT(act_half, alpha) * (other_half + beta).
// Which half is which depends on `act_first` (see below). Defaults
// alpha=1.0, beta=0.0 reproduce the plain SwiGLU/GeGLU behavior.
template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&, const float),
bool act_first, bool HAS_CLAMP>
__device__ __forceinline__ scalar_t compute(const scalar_t& x,
const scalar_t& y,
const float limit) {
const float limit,
const float alpha,
const float beta) {
if constexpr (act_first) {
scalar_t gate = x;
scalar_t up = y;
@@ -22,7 +31,9 @@ __device__ __forceinline__ scalar_t compute(const scalar_t& x,
gate = (scalar_t)fminf((float)gate, limit);
up = (scalar_t)fmaxf(fminf((float)up, limit), -limit);
}
return ACT_FN(gate) * up;
// act_first: gate is the activated half -> alpha applies to gate;
// beta is added to up (the non-activated half).
return ACT_FN(gate, alpha) * (scalar_t)((float)up + beta);
} else {
scalar_t gate = x;
scalar_t up = y;
@@ -30,55 +41,66 @@ __device__ __forceinline__ scalar_t compute(const scalar_t& x,
gate = (scalar_t)fmaxf(fminf((float)gate, limit), -limit);
up = (scalar_t)fminf((float)up, limit);
}
return gate * ACT_FN(up);
// !act_first: up is the activated half -> alpha applies to up;
// beta is added to gate (the non-activated half).
return (scalar_t)((float)gate + beta) * ACT_FN(up, alpha);
}
}
template <typename packed_t, packed_t (*PACKED_ACT_FN)(const packed_t&),
template <typename packed_t,
packed_t (*PACKED_ACT_FN)(const packed_t&, const float),
bool act_first, bool HAS_CLAMP>
__device__ __forceinline__ packed_t packed_compute(const packed_t& x,
const packed_t& y,
const float limit) {
const float limit,
const float alpha,
const float beta) {
if constexpr (act_first) {
packed_t gate = x;
packed_t up = y;
float2 u = cast_to_float2(up);
if constexpr (HAS_CLAMP) {
float2 g = cast_to_float2(gate);
float2 u = cast_to_float2(up);
g.x = fminf(g.x, limit);
g.y = fminf(g.y, limit);
u.x = fmaxf(fminf(u.x, limit), -limit);
u.y = fmaxf(fminf(u.y, limit), -limit);
gate = cast_to_packed<packed_t>(g);
up = cast_to_packed<packed_t>(u);
}
return packed_mul(PACKED_ACT_FN(gate), up);
// act_first: gate is the activated half -> alpha applies to gate;
// beta is added to up (the non-activated half).
u.x += beta;
u.y += beta;
return packed_mul(PACKED_ACT_FN(gate, alpha), cast_to_packed<packed_t>(u));
} else {
packed_t gate = x;
packed_t up = y;
float2 g = cast_to_float2(gate);
if constexpr (HAS_CLAMP) {
float2 g = cast_to_float2(gate);
float2 u = cast_to_float2(up);
g.x = fmaxf(fminf(g.x, limit), -limit);
g.y = fmaxf(fminf(g.y, limit), -limit);
u.x = fminf(u.x, limit);
u.y = fminf(u.y, limit);
gate = cast_to_packed<packed_t>(g);
up = cast_to_packed<packed_t>(u);
}
return packed_mul(gate, PACKED_ACT_FN(up));
// !act_first: up is the activated half -> alpha applies to up;
// beta is added to gate (the non-activated half).
g.x += beta;
g.y += beta;
return packed_mul(cast_to_packed<packed_t>(g), PACKED_ACT_FN(up, alpha));
}
}
// Activation and gating kernel template.
template <typename scalar_t, typename packed_t,
scalar_t (*ACT_FN)(const scalar_t&),
packed_t (*PACKED_ACT_FN)(const packed_t&), bool act_first,
bool use_vec, bool HAS_CLAMP, bool use_256b = false>
scalar_t (*ACT_FN)(const scalar_t&, const float),
packed_t (*PACKED_ACT_FN)(const packed_t&, const float),
bool act_first, bool use_vec, bool HAS_CLAMP, bool use_256b = false>
__global__ void act_and_mul_kernel(
scalar_t* __restrict__ out, // [..., d]
const scalar_t* __restrict__ input, // [..., 2, d]
const int d, const float limit) {
const int d, const float limit, const float alpha, const float beta) {
const scalar_t* x_ptr = input + blockIdx.x * 2 * d;
const scalar_t* y_ptr = x_ptr + d;
scalar_t* out_ptr = out + blockIdx.x * d;
@@ -105,7 +127,7 @@ __global__ void act_and_mul_kernel(
for (int j = 0; j < pvec_t::NUM_ELTS; j++) {
x.elts[j] =
packed_compute<packed_t, PACKED_ACT_FN, act_first, HAS_CLAMP>(
x.elts[j], y.elts[j], limit);
x.elts[j], y.elts[j], limit, alpha, beta);
}
if constexpr (use_256b) {
st256(x, &out_vec[i]);
@@ -118,29 +140,34 @@ __global__ void act_and_mul_kernel(
for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
const scalar_t x = VLLM_LDG(&x_ptr[idx]);
const scalar_t y = VLLM_LDG(&y_ptr[idx]);
out_ptr[idx] =
compute<scalar_t, ACT_FN, act_first, HAS_CLAMP>(x, y, limit);
out_ptr[idx] = compute<scalar_t, ACT_FN, act_first, HAS_CLAMP>(
x, y, limit, alpha, beta);
}
}
}
// Gated activations take an `alpha` argument that scales the sigmoid input
// (`x * sigmoid(alpha * x)`). alpha defaults to 1.0 at all call sites, which
// is exactly SiLU; only the clamp path (silu_and_mul_with_clamp) passes a
// non-default alpha. Activations that do not use alpha simply ignore it.
template <typename T>
__device__ __forceinline__ T silu_kernel(const T& x) {
// x * sigmoid(x)
return (T)(((float)x) / (1.0f + expf((float)-x)));
__device__ __forceinline__ T silu_kernel(const T& x, const float alpha) {
// x * sigmoid(alpha * x)
return (T)(((float)x) / (1.0f + expf((float)-x * alpha)));
}
template <typename packed_t>
__device__ __forceinline__ packed_t packed_silu_kernel(const packed_t& val) {
// x * sigmoid(x)
__device__ __forceinline__ packed_t packed_silu_kernel(const packed_t& val,
const float alpha) {
// x * sigmoid(alpha * x)
float2 fval = cast_to_float2(val);
fval.x = fval.x / (1.0f + expf(-fval.x));
fval.y = fval.y / (1.0f + expf(-fval.y));
fval.x = fval.x / (1.0f + expf(-fval.x * alpha));
fval.y = fval.y / (1.0f + expf(-fval.y * alpha));
return cast_to_packed<packed_t>(fval);
}
template <typename T>
__device__ __forceinline__ T gelu_kernel(const T& x) {
__device__ __forceinline__ T gelu_kernel(const T& x, const float /*alpha*/) {
// Equivalent to PyTorch GELU with 'none' approximation.
// Refer to:
// https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L36-L38
@@ -150,7 +177,8 @@ __device__ __forceinline__ T gelu_kernel(const T& x) {
}
template <typename packed_t>
__device__ __forceinline__ packed_t packed_gelu_kernel(const packed_t& val) {
__device__ __forceinline__ packed_t packed_gelu_kernel(const packed_t& val,
const float /*alpha*/) {
// Equivalent to PyTorch GELU with 'none' approximation.
// Refer to:
// https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L36-L38
@@ -162,7 +190,8 @@ __device__ __forceinline__ packed_t packed_gelu_kernel(const packed_t& val) {
}
template <typename T>
__device__ __forceinline__ T gelu_tanh_kernel(const T& x) {
__device__ __forceinline__ T gelu_tanh_kernel(const T& x,
const float /*alpha*/) {
// Equivalent to PyTorch GELU with 'tanh' approximation.
// Refer to:
// https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L25-L30
@@ -176,7 +205,7 @@ __device__ __forceinline__ T gelu_tanh_kernel(const T& x) {
template <typename packed_t>
__device__ __forceinline__ packed_t
packed_gelu_tanh_kernel(const packed_t& val) {
packed_gelu_tanh_kernel(const packed_t& val, const float /*alpha*/) {
// Equivalent to PyTorch GELU with 'tanh' approximation.
// Refer to:
// https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L25-L30
@@ -202,7 +231,7 @@ packed_gelu_tanh_kernel(const packed_t& val) {
// clamped (max only) and up input is clamped (both sides) before the
// activation function is applied.
#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST, \
HAS_CLAMP, LIMIT) \
HAS_CLAMP, LIMIT, ALPHA, BETA) \
auto dtype = input.scalar_type(); \
int d = input.size(-1) / 2; \
int64_t num_tokens = input.numel() / input.size(-1); \
@@ -230,7 +259,7 @@ packed_gelu_tanh_kernel(const packed_t& val) {
PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
ACT_FIRST, true, HAS_CLAMP, true><<<grid, block, 0, stream>>>( \
out.mutable_data_ptr<scalar_t>(), \
input.const_data_ptr<scalar_t>(), d, LIMIT); \
input.const_data_ptr<scalar_t>(), d, LIMIT, ALPHA, BETA); \
}); \
} else { \
VLLM_STABLE_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] { \
@@ -240,7 +269,7 @@ packed_gelu_tanh_kernel(const packed_t& val) {
PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
ACT_FIRST, true, HAS_CLAMP, false><<<grid, block, 0, stream>>>( \
out.mutable_data_ptr<scalar_t>(), \
input.const_data_ptr<scalar_t>(), d, LIMIT); \
input.const_data_ptr<scalar_t>(), d, LIMIT, ALPHA, BETA); \
}); \
} \
} else { \
@@ -252,7 +281,7 @@ packed_gelu_tanh_kernel(const packed_t& val) {
PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
ACT_FIRST, false, HAS_CLAMP><<<grid, block, 0, stream>>>( \
out.mutable_data_ptr<scalar_t>(), input.const_data_ptr<scalar_t>(), \
d, LIMIT); \
d, LIMIT, ALPHA, BETA); \
}); \
}
@@ -260,14 +289,18 @@ void silu_and_mul(torch::stable::Tensor& out, // [..., d]
torch::stable::Tensor& input) // [..., 2 * d]
{
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
true, false, 0.0f);
true, false, 0.0f, 1.0f, 0.0f);
}
void silu_and_mul_clamp(torch::stable::Tensor& out, // [..., d]
torch::stable::Tensor& input, // [..., 2 * d]
double limit) {
double limit, double alpha, double beta) {
// out = (gate.clamp(max=limit) * sigmoid(alpha * gate.clamp(max=limit)))
// * (up.clamp(+-limit) + beta)
// alpha=1.0, beta=0.0 reduce this to silu(gate) * up.
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
true, true, (float)limit);
true, true, (float)limit, (float)alpha,
(float)beta);
}
void mul_and_silu(torch::stable::Tensor& out, // [..., d]
@@ -276,21 +309,22 @@ void mul_and_silu(torch::stable::Tensor& out, // [..., d]
// The difference between mul_and_silu and silu_and_mul is that mul_and_silu
// applies the silu to the latter half of the input.
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
false, false, 0.0f);
false, false, 0.0f, 1.0f, 0.0f);
}
void gelu_and_mul(torch::stable::Tensor& out, // [..., d]
torch::stable::Tensor& input) // [..., 2 * d]
{
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel, vllm::packed_gelu_kernel,
true, false, 0.0f);
true, false, 0.0f, 1.0f, 0.0f);
}
void gelu_tanh_and_mul(torch::stable::Tensor& out, // [..., d]
torch::stable::Tensor& input) // [..., 2 * d]
{
LAUNCH_ACTIVATION_GATE_KERNEL(
vllm::gelu_tanh_kernel, vllm::packed_gelu_tanh_kernel, true, false, 0.0f);
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_tanh_kernel,
vllm::packed_gelu_tanh_kernel, true, false,
0.0f, 1.0f, 0.0f);
}
namespace vllm {
@@ -2,7 +2,7 @@
#include <torch/csrc/stable/tensor.h>
#include "broadcast_load_epilogue_c2x.hpp"
#include "cutlass_extensions/epilogue/broadcast_load_epilogue_c2x.hpp"
/*
This file defines custom epilogues for fusing channel scales, token scales,
@@ -0,0 +1,615 @@
/*
* SPDX-License-Identifier: Apache-2.0
* SPDX-FileCopyrightText: Copyright contributors to the vLLM project
*
* Horizontally-fused MiniMax-M3 attention pre-processing kernel.
*
* Replaces the per-token Python sequence in
* ``MiniMaxM3SparseAttention.forward`` / ``MiniMaxM3Attention.forward``:
*
* q = q_norm(q); k = k_norm(k); q, k = rotary_emb(pos, q, k)
* index_q = index_q_norm(index_q); index_k = index_k_norm(index_k)
* index_q, index_k = rotary_emb(pos, index_q, index_k)
* _insert_kv(k, v, index_k)
*
* All branches share head_dim=128 and the *same* partial-NeoX RoPE table
* (``rotary_dim`` rotated, the trailing dims pass through). The four norms
* are Gemma-style RMSNorm (``x * rsqrt(mean(x^2)+eps) * (1 + weight)``) with
* independent weights.
*
* Everything lives in a single fused ``qkv`` tensor. The sparse layer's
* fused projection (MinimaxM3QKVParallelLinearWithIndexer) emits, per token::
*
* [ q | k | v | index_q | index_k ] (the "5 results")
*
* while the dense layer emits just ``[ q | k | v ]``. The kernel reads the
* index branch straight out of that packed row -- no separate index tensors.
*
* One kernel, one grid; each warp owns one (token, head-slot) pair. Slot
* enumeration per token:
* [0, nq) Q heads -> norm(q_w) + RoPE, write
* qkv [nq, nq+nkv) K heads -> norm(k_w) + RoPE, write
* qkv
* (+ insert into key cache)
* [nq+nkv, nq+2*nkv) V heads -> insert into value cache
* IQ heads (niq) -> norm(iq_w) + RoPE, write iq
* IK (1) -> norm(ik_w) + RoPE
* (+ insert into index cache)
*
* The IQ/IK warps address the index_q/index_k sub-blocks *inside* qkv at the
* fixed physical offsets (nq+2*nkv)*128 and (nq+2*nkv+niq)*128.
*
* Dense vs sparse is a compile-time choice via the ``kIsSparse``/``kInsertKV``
* template bools (3 instantiations: dense <false,false>, sparse-profiling
* <true,false>, sparse-serving <true,true>), so the index slots, the V slots
* and the cache inserts fold away entirely on paths that don't use them. The
* dense layer passes no caches/index: norm+RoPE happens in place and the
* generic ``Attention`` layer owns the cache write.
*
* Q/K and (sparse) index_q/index_k are all rewritten in place inside the fused
* ``qkv`` tensor. Caches (bf16) are scatter-written by slot.
*/
#include <cmath>
#include <cuda_runtime.h>
#include <type_traits>
#include "torch_utils.h"
#include "../cuda_compat.h"
#include "../type_convert.cuh"
#include "dispatch_utils.h"
#ifndef FINAL_MASK
#ifdef USE_ROCM
#define FINAL_MASK 0xffffffffffffffffULL
#else
#define FINAL_MASK 0xffffffffu
#endif
#endif
namespace vllm {
namespace minimax_m3_fused_ops {
namespace {
inline int getSMVersion() {
auto* props = get_device_prop();
return props->major * 10 + props->minor;
}
} // namespace
// ────────────────────────────────────────────────────────────────────────────
// Constants (hard-coded for MiniMax-M3-preview).
// ────────────────────────────────────────────────────────────────────────────
constexpr int kHeadDim = 128;
constexpr int kNumLanes = 32;
constexpr int kElemsPerLane = kHeadDim / kNumLanes; // 4
// ────────────────────────────────────────────────────────────────────────────
// Helpers
// ────────────────────────────────────────────────────────────────────────────
__device__ __forceinline__ float warpReduceSum(float val) {
#pragma unroll
for (int mask = 16; mask > 0; mask >>= 1) {
val += __shfl_xor_sync(FINAL_MASK, val, mask, 32);
}
return val;
}
// Gemma RMSNorm over the full head (no-op when ``weight == nullptr``) followed
// by partial NeoX RoPE on the leading ``rotary_dim`` dims, all in fp32. Each
// lane owns ``kElemsPerLane`` contiguous dims [laneId*4, laneId*4+4).
template <typename scalar_t>
__device__ __forceinline__ void normAndRope(
float (&elems)[kElemsPerLane], int const laneId, float const eps,
scalar_t const* __restrict__ weight, // [kHeadDim] or nullptr (no norm)
bool const do_rope, int const rotary_dim,
scalar_t const* __restrict__ cos_ptr, // cos_sin_cache + pos*rotary_dim
bool const apply_norm) {
// ── Gemma RMSNorm: x * rsqrt(mean(x^2)+eps) * (1 + w) ──────────────────
if (apply_norm) {
float sumsq = 0.0f;
#pragma unroll
for (int i = 0; i < kElemsPerLane; i++) sumsq += elems[i] * elems[i];
sumsq = warpReduceSum(sumsq);
float const rms_rcp = rsqrtf(sumsq / static_cast<float>(kHeadDim) + eps);
#pragma unroll
for (int i = 0; i < kElemsPerLane; i++) {
int const dim = laneId * kElemsPerLane + i;
float const w = 1.0f + static_cast<float>(weight[dim]);
elems[i] = elems[i] * rms_rcp * w;
}
}
// ── Partial NeoX RoPE on dims [0, rotary_dim) ──────────────────────────
// half = rotary_dim/2. Pair (i, i+half) for i in [0, half). Lane L owns
// dims [4L, 4L+4); since half is a multiple of 4, a lane lies wholly in the
// first half (own=x[i]) or second half (own=x[i+half]); its partner lives
// ``half/4`` lanes away (XOR with that distance).
if (do_rope) {
int const half = rotary_dim / 2;
int const dim0 = laneId * kElemsPerLane;
bool const in_rope = dim0 < rotary_dim;
int const lane_xor = half / kElemsPerLane; // partner-lane distance
float partner[kElemsPerLane];
#pragma unroll
for (int i = 0; i < kElemsPerLane; i++) {
partner[i] = __shfl_xor_sync(FINAL_MASK, elems[i], lane_xor, 32);
}
if (in_rope) {
bool const first_half = dim0 < half;
int const i_base = first_half ? dim0 : (dim0 - half); // cos/sin index
scalar_t const* sin_ptr = cos_ptr + half;
#pragma unroll
for (int i = 0; i < kElemsPerLane; i++) {
float const c = static_cast<float>(cos_ptr[i_base + i]);
float const s = static_cast<float>(sin_ptr[i_base + i]);
if (first_half) {
elems[i] = elems[i] * c - partner[i] * s;
} else {
elems[i] = elems[i] * c + partner[i] * s;
}
}
}
}
}
// Load 4 contiguous bf16 -> 4 fp32 registers.
template <typename scalar_t>
__device__ __forceinline__ void loadElems(scalar_t const* __restrict__ src,
float (&elems)[kElemsPerLane]) {
using Converter = vllm::_typeConvert<scalar_t>;
uint2 v = *reinterpret_cast<uint2 const*>(src);
auto const* p =
reinterpret_cast<typename Converter::packed_hip_type const*>(&v);
#pragma unroll
for (int i = 0; i < kElemsPerLane / 2; i++) {
float2 f2 = Converter::convert(p[i]);
elems[2 * i] = f2.x;
elems[2 * i + 1] = f2.y;
}
}
// Store 4 fp32 registers -> 4 contiguous bf16.
template <typename scalar_t>
__device__ __forceinline__ void storeElems(
scalar_t* __restrict__ dst, float const (&elems)[kElemsPerLane]) {
using Converter = vllm::_typeConvert<scalar_t>;
uint2 v;
auto* p = reinterpret_cast<typename Converter::packed_hip_type*>(&v);
#pragma unroll
for (int i = 0; i < kElemsPerLane / 2; i++) {
p[i] = Converter::convert(make_float2(elems[2 * i], elems[2 * i + 1]));
}
*reinterpret_cast<uint2*>(dst) = v;
}
// ────────────────────────────────────────────────────────────────────────────
// Kernel
// ────────────────────────────────────────────────────────────────────────────
// Grid: 1D, ceil(num_tokens * slots_per_token / warps_per_block).
// Each warp = one (token, slot).
//
// `kIsSparse` and `kInsertKV` are compile-time template bools, so all the
// branch decisions that distinguish the dense layer from the sparse layer
// (index slots, KV/index inserts, V slots) fold away per instantiation.
// Three instantiations are built: dense <false,false>, sparse-profiling
// <true,false> and sparse-serving <true,true>. Slots per token:
// Q : nq (always — norm+RoPE)
// K : nkv (always — norm+RoPE; +K-cache insert)
// V : nkv only if kInsertKV (V-cache insert; no warps in dense)
// IQ: niq only if kIsSparse (norm+RoPE)
// IK: 1 only if kIsSparse (norm+RoPE; +index-cache insert)
template <typename scalar_t, bool kIsSparse, bool kInsertKV>
__global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
scalar_t* __restrict__ qkv, // [N, qkv_row] in/out (packs index if sparse)
scalar_t* __restrict__ q_out, // [N, nq*128] contiguous, or nullptr
scalar_t* __restrict__ index_q_out, // [N, niq*128] contiguous, or nullptr
scalar_t const* __restrict__ q_norm_w,
scalar_t const* __restrict__ k_norm_w,
scalar_t const* __restrict__ iq_norm_w,
scalar_t const* __restrict__ ik_norm_w,
scalar_t const* __restrict__ cos_sin_cache, // [max_pos, rotary_dim]
int64_t const* __restrict__ positions, // [N] i64
int64_t const* __restrict__ slot_mapping, // [N] i64 or nullptr
scalar_t* __restrict__ kv_cache, // [nb,2,bs,nkv,128] or nullptr
scalar_t* __restrict__ index_cache, // [nb*bs, 128] or nullptr
float const eps, int const rotary_dim, int const num_tokens, int const nq,
int const nkv, int const niq, int const block_size,
// kv_cache strides (in elements) for logical shape [nb, 2, bs, nkv, 128].
// The head_dim (last) dim is always innermost-contiguous (stride 1), so the
// NHD/HND layout choice is fully captured by these four strides: NHD keeps
// s_token < s_head, HND swaps them. dim_base addresses head_dim directly.
int64_t const kv_s_block, int64_t const kv_s_kv, int64_t const kv_s_token,
int64_t const kv_s_head) {
#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
// _typeConvert<BFloat16> is unavailable on pre-Ampere; the M3 kernel only
// runs with bf16/fp16 inputs in practice. Discard the bf16 body there.
if constexpr (std::is_same_v<scalar_t, c10::BFloat16>) {
return;
} else {
#endif
int const warpsPerBlock = blockDim.x / 32;
int const laneId = threadIdx.x % 32;
int const globalWarpIdx = blockIdx.x * warpsPerBlock + (threadIdx.x / 32);
// Slot layout (compile-time gated: dense has neither V nor index slots).
int const v_slots = kInsertKV ? nkv : 0;
int const idx_slots = kIsSparse ? niq + 1 : 0;
int const slots_per_token = nq + nkv + v_slots + idx_slots;
int const tokenIdx = globalWarpIdx / slots_per_token;
int const slot = globalWarpIdx % slots_per_token;
if (tokenIdx >= num_tokens) return;
// Slot boundaries.
int const k_begin = nq;
int const v_begin = nq + nkv; // valid only when kInsertKV
int const iq_begin = nq + nkv + v_slots; // index block start
int const ik_slot = iq_begin + niq; // valid only when kIsSparse
bool const isQ = slot < k_begin;
bool const isK = slot >= k_begin && slot < v_begin;
bool isV = false;
if constexpr (kInsertKV) isV = slot >= v_begin && slot < v_begin + nkv;
bool isIQ = false, isIK = false;
if constexpr (kIsSparse) {
isIQ = slot >= iq_begin && slot < ik_slot;
isIK = slot == ik_slot;
}
int const dim_base = laneId * kElemsPerLane;
// Physical row width of qkv: the dense layer packs [q|k|v]; the sparse
// layer additionally packs [index_q (niq heads) | index_k (1 head)].
int const qkv_row = (nq + 2 * nkv + (kIsSparse ? (niq + 1) : 0)) * kHeadDim;
// ── Resolve source pointer + per-branch parameters. ────────────────────
scalar_t* row_ptr = nullptr; // in-place output location
scalar_t const* norm_w = nullptr; // nullptr -> skip norm (V)
bool do_rope = true;
int head = 0; // kv head index for inserts
if (isQ) {
row_ptr =
qkv + static_cast<int64_t>(tokenIdx) * qkv_row + slot * kHeadDim;
norm_w = q_norm_w;
} else if (isK) {
head = slot - k_begin;
row_ptr =
qkv + static_cast<int64_t>(tokenIdx) * qkv_row + slot * kHeadDim;
norm_w = k_norm_w;
} else if (isV) {
// qkv V section starts at slot index (nq + nkv): slot * kHeadDim is the
// correct in-tensor offset.
head = slot - v_begin;
row_ptr =
qkv + static_cast<int64_t>(tokenIdx) * qkv_row + slot * kHeadDim;
norm_w = nullptr; // V: no norm, no rope
do_rope = false;
} else if (isIQ) {
// index_q sub-block lives at physical offset (nq+2*nkv)*128 in qkv.
int const ih = slot - iq_begin;
row_ptr = qkv + static_cast<int64_t>(tokenIdx) * qkv_row +
(nq + 2 * nkv + ih) * kHeadDim;
norm_w = iq_norm_w;
} else { // isIK -- single shared index key at (nq+2*nkv+niq)*128.
row_ptr = qkv + static_cast<int64_t>(tokenIdx) * qkv_row +
(nq + 2 * nkv + niq) * kHeadDim;
norm_w = ik_norm_w;
}
// Store destination. Q and index_q are gathered into dedicated contiguous
// output buffers (when provided) so the downstream SM100 sparse kernel's
// flat TMA descriptor can address them as [tokens*heads, head_dim]; this
// folds the de-interleaving into the store the kernel already does, instead
// of a separate q.contiguous() copy. Everything else stays in place.
scalar_t* store_ptr = row_ptr;
if (isQ && q_out != nullptr) {
store_ptr = q_out + static_cast<int64_t>(tokenIdx) * nq * kHeadDim +
slot * kHeadDim;
} else if (isIQ && index_q_out != nullptr) {
store_ptr = index_q_out +
static_cast<int64_t>(tokenIdx) * niq * kHeadDim +
(slot - iq_begin) * kHeadDim;
}
// PDL: wait for the predecessor kernel (the qkv-projection GEMM that
// produces ``qkv``) to finish before touching any global memory. No-op
// when PDL is not enabled on the launch. The CUDA runtime wrapper emits
// the griddepcontrol.wait PTX with the required memory clobber internally.
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
cudaGridDependencySynchronize();
#endif
// ── Load -> norm+rope (fp32) -> store back in place. ───────────────────
float elems[kElemsPerLane];
loadElems<scalar_t>(row_ptr + dim_base, elems);
if (!isV) {
int64_t const pos = positions[tokenIdx];
scalar_t const* cos_ptr = cos_sin_cache + pos * rotary_dim;
normAndRope<scalar_t>(elems, laneId, eps, norm_w, do_rope, rotary_dim,
cos_ptr, /*apply_norm=*/norm_w != nullptr);
storeElems<scalar_t>(store_ptr + dim_base, elems);
}
// ── Cache inserts (sparse serving only). ───────────────────────────────
if constexpr (kInsertKV) {
// Guard (not early-return) so every thread reaches the PDL trigger below.
int64_t const sm = (isK || isV || isIK) ? slot_mapping[tokenIdx] : -1;
if (sm >= 0) { // skip padded / unscheduled tokens
if (isIK) {
scalar_t* dst = index_cache + sm * kHeadDim + dim_base;
storeElems<scalar_t>(dst, elems);
} else if (isK || isV) {
// kv_cache logical shape [num_blocks, 2, block_size, nkv, head_dim].
// Paging is logical (block = sm/block_size, token = sm%block_size);
// the physical NHD/HND layout is honoured via the passed strides.
int64_t const b = sm / block_size;
int64_t const t = sm % block_size;
int const kv = isK ? 0 : 1;
int64_t const off =
b * kv_s_block + kv * kv_s_kv + t * kv_s_token + head * kv_s_head;
storeElems<scalar_t>(kv_cache + off + dim_base, elems);
}
}
}
// PDL: signal that this kernel is done so a dependent successor may launch
// early. No-op when PDL is not enabled on the launch.
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
cudaTriggerProgrammaticLaunchCompletion();
#endif
#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
}
#endif
}
// ────────────────────────────────────────────────────────────────────────────
// Launch wrapper
// ────────────────────────────────────────────────────────────────────────────
template <typename scalar_t>
void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
scalar_t const* q_norm_w, scalar_t const* k_norm_w,
scalar_t const* iq_norm_w, scalar_t const* ik_norm_w,
scalar_t const* cos_sin_cache,
int64_t const* positions, int64_t const* slot_mapping,
scalar_t* kv_cache, scalar_t* index_cache,
float const eps, int const rotary_dim,
int const num_tokens, int const nq, int const nkv,
int const niq, int const block_size,
int64_t const kv_s_block, int64_t const kv_s_kv,
int64_t const kv_s_token, int64_t const kv_s_head,
bool const has_index, bool const insert_kv,
cudaStream_t stream) {
// Slot count must match the kernel's compile-time gating.
int const v_slots = insert_kv ? nkv : 0;
int const idx_slots = has_index ? niq + 1 : 0;
int const slots_per_token = nq + nkv + v_slots + idx_slots;
constexpr int kBlockSize = 256;
constexpr int kWarpsPerBlock = kBlockSize / 32;
int64_t const total_warps =
static_cast<int64_t>(num_tokens) * slots_per_token;
int const grid =
static_cast<int>((total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock);
if (grid == 0) return;
#ifndef USE_ROCM
// PDL: enable programmatic stream serialization whenever the hardware
// supports it (SM90+). On pre-Hopper GPUs the attribute is unavailable, so
// leave numAttrs = 0 and launch as a regular kernel via cudaLaunchKernelEx.
static int const sm_version = getSMVersion();
cudaLaunchConfig_t config;
config.gridDim = dim3(grid);
config.blockDim = dim3(kBlockSize);
config.dynamicSmemBytes = 0;
config.stream = stream;
cudaLaunchAttribute attrs[1];
attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
attrs[0].val.programmaticStreamSerializationAllowed = 1;
config.attrs = attrs;
config.numAttrs = (sm_version >= 90) ? 1 : 0;
#define LAUNCH(IS_SPARSE, INSERT) \
cudaLaunchKernelEx( \
&config, \
fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, IS_SPARSE, INSERT>, \
qkv, q_out, index_q_out, q_norm_w, k_norm_w, iq_norm_w, ik_norm_w, \
cos_sin_cache, positions, slot_mapping, kv_cache, index_cache, eps, \
rotary_dim, num_tokens, nq, nkv, niq, block_size, kv_s_block, kv_s_kv, \
kv_s_token, kv_s_head)
#else
// ROCm: standard kernel launch syntax (no PDL/stream serialization).
// clang-format off
#define LAUNCH(IS_SPARSE, INSERT) \
fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, IS_SPARSE, INSERT> \
<<<grid, kBlockSize, 0, stream>>>( \
qkv, q_out, index_q_out, q_norm_w, k_norm_w, iq_norm_w, \
ik_norm_w, cos_sin_cache, positions, slot_mapping, kv_cache, \
index_cache, eps, rotary_dim, num_tokens, nq, nkv, niq, \
block_size, kv_s_block, kv_s_kv, kv_s_token, kv_s_head)
// clang-format on
#endif
if (has_index) {
if (insert_kv) {
LAUNCH(true, true); // sparse serving
} else {
LAUNCH(true, false); // sparse profiling
}
} else {
// Dense layer: never has an index branch and never inserts here (the
// generic Attention layer owns the KV insert).
LAUNCH(false, false);
}
#undef LAUNCH
}
} // namespace minimax_m3_fused_ops
} // namespace vllm
// ────────────────────────────────────────────────────────────────────────────
// Torch op wrapper
// ────────────────────────────────────────────────────────────────────────────
void fused_minimax_m3_qknorm_rope_kv_insert(
torch::stable::Tensor& qkv, // [N, qkv_row] (packs index if sparse)
torch::stable::Tensor const& q_norm_weight, // [128]
torch::stable::Tensor const& k_norm_weight, // [128]
torch::stable::Tensor const& cos_sin_cache, // [max_pos, rotary_dim]
torch::stable::Tensor const& positions, // [N] i64
int64_t num_heads, int64_t num_kv_heads, int64_t rotary_dim, double eps,
std::optional<torch::stable::Tensor> index_q_norm_weight, // [128]
std::optional<torch::stable::Tensor> index_k_norm_weight, // [128]
int64_t num_index_heads, // niq; 0 => dense
std::optional<torch::stable::Tensor> slot_mapping, // [N] i64
std::optional<torch::stable::Tensor> kv_cache, // [nb,2,bs,nkv,128]
std::optional<torch::stable::Tensor> index_cache, // [nb,bs,128]
int64_t block_size,
std::optional<torch::stable::Tensor> q_out, // [N, nq*128] contiguous
std::optional<torch::stable::Tensor>
index_q_out) { // [N, niq*128] contiguous
STD_TORCH_CHECK(qkv.is_cuda() && qkv.is_contiguous(),
"qkv must be contiguous CUDA");
STD_TORCH_CHECK(
positions.is_cuda() &&
positions.scalar_type() == torch::headeronly::ScalarType::Long,
"positions must be int64 CUDA");
STD_TORCH_CHECK(cos_sin_cache.is_cuda() && cos_sin_cache.is_contiguous(),
"cos_sin_cache must be contiguous CUDA");
STD_TORCH_CHECK(cos_sin_cache.scalar_type() == qkv.scalar_type(),
"cos_sin_cache dtype must match qkv");
STD_TORCH_CHECK(
cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == rotary_dim,
"cos_sin_cache shape [max_pos, rotary_dim]");
STD_TORCH_CHECK(q_norm_weight.scalar_type() == qkv.scalar_type() &&
k_norm_weight.scalar_type() == qkv.scalar_type(),
"q/k norm weight dtype must match qkv");
STD_TORCH_CHECK(
q_norm_weight.numel() == vllm::minimax_m3_fused_ops::kHeadDim &&
k_norm_weight.numel() == vllm::minimax_m3_fused_ops::kHeadDim,
"q/k norm weight must have 128 elements");
STD_TORCH_CHECK(rotary_dim > 0 && rotary_dim % 8 == 0 &&
rotary_dim <= vllm::minimax_m3_fused_ops::kHeadDim,
"rotary_dim must be a positive multiple of 8 and <= 128");
int const num_tokens = static_cast<int>(qkv.size(0));
int const nq = static_cast<int>(num_heads);
int const nkv = static_cast<int>(num_kv_heads);
int const niq = static_cast<int>(num_index_heads);
// The sparse layer packs the index branch ([index_q (niq heads) | index_k
// (1 head)]) right after [q|k|v] in the same row; the dense layer does not.
bool const has_index = niq > 0;
bool const insert_kv = kv_cache.has_value();
int const kHeadDim = vllm::minimax_m3_fused_ops::kHeadDim;
int const expected_row =
(nq + 2 * nkv + (has_index ? niq + 1 : 0)) * kHeadDim;
STD_TORCH_CHECK(qkv.size(1) == expected_row,
"qkv last dim must be (num_heads + 2*num_kv_heads"
" + num_index_heads + 1) * 128 for sparse, "
"(num_heads + 2*num_kv_heads) * 128 for dense");
// Only the sparse layer inserts here (dense lets the generic Attention layer
// own the KV write); there is no dense+insert kernel instantiation.
STD_TORCH_CHECK(
!insert_kv || has_index,
"insert mode (kv_cache) requires the index branch (sparse layer)");
if (has_index) {
STD_TORCH_CHECK(
index_q_norm_weight.has_value() && index_k_norm_weight.has_value(),
"index branch requires both index norm weights");
STD_TORCH_CHECK(index_q_norm_weight->scalar_type() == qkv.scalar_type() &&
index_k_norm_weight->scalar_type() == qkv.scalar_type(),
"index norm weights dtype must match qkv");
STD_TORCH_CHECK(index_q_norm_weight->numel() == kHeadDim &&
index_k_norm_weight->numel() == kHeadDim,
"index norm weights must have 128 elements");
}
// kv_cache strides (logical shape [nb, 2, bs, nkv, head_dim]). Read straight
// off the tensor so the kernel honours whatever physical layout the attention
// backend allocated (NHD: stride order (0,1,2,3,4); HND: (0,1,3,2,4)). No new
// op argument is needed -- the strides ride along with the tensor itself.
int64_t kv_s_block = 0, kv_s_kv = 0, kv_s_token = 0, kv_s_head = 0;
if (insert_kv) {
STD_TORCH_CHECK(
slot_mapping.has_value() &&
slot_mapping->scalar_type() == torch::headeronly::ScalarType::Long,
"insert mode requires int64 slot_mapping");
STD_TORCH_CHECK(kv_cache->scalar_type() == qkv.scalar_type(),
"kv_cache dtype must match qkv (bf16 cache only)");
STD_TORCH_CHECK(index_cache.has_value() &&
index_cache->scalar_type() == qkv.scalar_type(),
"insert mode requires matching index_cache");
STD_TORCH_CHECK(kv_cache->dim() == 5 && kv_cache->stride(4) == 1,
"kv_cache must be [nb,2,bs,nkv,head_dim] with contiguous "
"head_dim (stride(4)==1)");
kv_s_block = kv_cache->stride(0);
kv_s_kv = kv_cache->stride(1);
kv_s_token = kv_cache->stride(2);
kv_s_head = kv_cache->stride(3);
}
// Optional contiguous gather targets: when given, the normed/roped q (and
// index_q) are written here instead of in place, so callers avoid a separate
// .contiguous() copy. index_q_out only makes sense on the sparse path.
if (q_out.has_value()) {
STD_TORCH_CHECK(
q_out->is_cuda() && q_out->is_contiguous() &&
q_out->scalar_type() == qkv.scalar_type(),
"q_out must be a contiguous CUDA tensor matching qkv dtype");
STD_TORCH_CHECK(
q_out->numel() == static_cast<int64_t>(num_tokens) * nq * kHeadDim,
"q_out must have num_tokens * num_heads * 128 elements");
}
if (index_q_out.has_value()) {
STD_TORCH_CHECK(
has_index,
"index_q_out requires the index branch (num_index_heads > 0)");
STD_TORCH_CHECK(
index_q_out->is_cuda() && index_q_out->is_contiguous() &&
index_q_out->scalar_type() == qkv.scalar_type(),
"index_q_out must be a contiguous CUDA tensor matching qkv dtype");
STD_TORCH_CHECK(index_q_out->numel() ==
static_cast<int64_t>(num_tokens) * niq * kHeadDim,
"index_q_out must have num_tokens * num_index_heads * 128 "
"elements");
}
const torch::stable::accelerator::DeviceGuard device_guard(
qkv.get_device_index());
auto stream = get_current_cuda_stream(qkv.get_device_index());
VLLM_STABLE_DISPATCH_HALF_TYPES(
qkv.scalar_type(), "fused_minimax_m3_qknorm_rope_kv_insert", [&] {
using st = scalar_t;
vllm::minimax_m3_fused_ops::launchFusedMiniMaxM3<st>(
reinterpret_cast<st*>(qkv.data_ptr()),
q_out.has_value() ? reinterpret_cast<st*>(q_out->data_ptr())
: nullptr,
index_q_out.has_value()
? reinterpret_cast<st*>(index_q_out->data_ptr())
: nullptr,
reinterpret_cast<st const*>(q_norm_weight.data_ptr()),
reinterpret_cast<st const*>(k_norm_weight.data_ptr()),
has_index
? reinterpret_cast<st const*>(index_q_norm_weight->data_ptr())
: nullptr,
has_index
? reinterpret_cast<st const*>(index_k_norm_weight->data_ptr())
: nullptr,
reinterpret_cast<st const*>(cos_sin_cache.data_ptr()),
reinterpret_cast<int64_t const*>(positions.data_ptr()),
insert_kv
? reinterpret_cast<int64_t const*>(slot_mapping->data_ptr())
: nullptr,
insert_kv ? reinterpret_cast<st*>(kv_cache->data_ptr()) : nullptr,
(insert_kv && has_index)
? reinterpret_cast<st*>(index_cache->data_ptr())
: nullptr,
static_cast<float>(eps), static_cast<int>(rotary_dim), num_tokens,
nq, nkv, niq, static_cast<int>(block_size), kv_s_block, kv_s_kv,
kv_s_token, kv_s_head, has_index, insert_kv, stream);
});
}
@@ -20,7 +20,7 @@
#include "torch_utils.h"
#include "async_util.cuh"
#include "../async_util.cuh"
#include "../cuda_compat.h"
#include "../type_convert.cuh"
#include "dispatch_utils.h"
+19 -1
View File
@@ -231,6 +231,23 @@ void fused_qk_norm_rope(torch::stable::Tensor& qkv, int64_t num_heads_q,
torch::stable::Tensor& position_ids,
int64_t forced_token_heads_per_warp);
// Horizontally-fused MiniMax-M3 QK-norm + partial NeoX RoPE (+ optional KV /
// index-cache insert). Dense layer: norm+RoPE only; sparse layer: also packs
// the index branch and scatters k/v/index_k into their paged caches.
void fused_minimax_m3_qknorm_rope_kv_insert(
torch::stable::Tensor& qkv, torch::stable::Tensor const& q_norm_weight,
torch::stable::Tensor const& k_norm_weight,
torch::stable::Tensor const& cos_sin_cache,
torch::stable::Tensor const& positions, int64_t num_heads,
int64_t num_kv_heads, int64_t rotary_dim, double eps,
std::optional<torch::stable::Tensor> index_q_norm_weight,
std::optional<torch::stable::Tensor> index_k_norm_weight,
int64_t num_index_heads, std::optional<torch::stable::Tensor> slot_mapping,
std::optional<torch::stable::Tensor> kv_cache,
std::optional<torch::stable::Tensor> index_cache, int64_t block_size,
std::optional<torch::stable::Tensor> q_out,
std::optional<torch::stable::Tensor> index_q_out);
// Sampler kernels (shared CUDA/ROCm)
void apply_repetition_penalties_(
torch::stable::Tensor& logits, const torch::stable::Tensor& prompt_mask,
@@ -276,7 +293,8 @@ void selective_scan_fwd(
// Activation kernels (shared CUDA/ROCm)
void silu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
void silu_and_mul_clamp(torch::stable::Tensor& out,
torch::stable::Tensor& input, double limit);
torch::stable::Tensor& input, double limit,
double alpha = 1.0, double beta = 0.0);
void mul_and_silu(torch::stable::Tensor& out, torch::stable::Tensor& input);
void gelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
void gelu_tanh_and_mul(torch::stable::Tensor& out,
@@ -25,7 +25,7 @@
#include <cuda_fp8.h>
#include "cuda_utils.h"
#include "libtorch_stable/launch_bounds_utils.h"
#include "launch_bounds_utils.h"
// Define before including nvfp4_utils.cuh so the header
// can use this macro during compilation.
@@ -34,7 +34,7 @@
static_assert(CVT_FP4_ELTS_PER_THREAD == 16,
"MXFP4 experts quant requires PACK16 mode (CUDA >= 12.9)");
#include "libtorch_stable/launch_bounds_utils.h"
#include "launch_bounds_utils.h"
namespace vllm {
@@ -26,7 +26,7 @@
#include "cuda_utils.h"
#include "nvfp4_utils.cuh"
#include "libtorch_stable/launch_bounds_utils.h"
#include "launch_bounds_utils.h"
namespace vllm {
@@ -26,7 +26,7 @@
#include "../../cuda_vec_utils.cuh"
#include "cuda_utils.h"
#include "libtorch_stable/launch_bounds_utils.h"
#include "launch_bounds_utils.h"
// Define before including nvfp4_utils.cuh so the header
// can use this macro during compilation.
@@ -7,11 +7,14 @@
#include <torch/csrc/stable/ops.h>
#include "ggml-common.h"
#include "vecdotq.cuh"
#include "dequantize.cuh"
#include "mmvq.cuh"
#include "mmq.cuh"
// NOTE: These headers are intentionally kept in csrc/quantization/gguf/ (not
// moved to libtorch_stable) to avoid unnecessary reformatting that would break
// git rename detection and pollute blame history.
#include "../../../quantization/gguf/ggml-common.h"
#include "../../../quantization/gguf/vecdotq.cuh"
#include "../../../quantization/gguf/dequantize.cuh"
#include "../../../quantization/gguf/mmvq.cuh"
#include "../../../quantization/gguf/mmq.cuh"
#include "moe.cuh"
#include "moe_vec.cuh"
+1 -1
View File
@@ -7,7 +7,7 @@
#include "torch_utils.h"
#ifndef USE_ROCM
#include "persistent_topk.cuh"
#include "../persistent_topk.cuh"
#endif
namespace {
+17 -2
View File
@@ -337,6 +337,17 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
"bool is_neox, Tensor position_ids, "
"int forced_token_heads_per_warp=-1) -> ()");
// Horizontally-fused MiniMax-M3 QK-norm + partial NeoX RoPE + KV-insert.
ops.def(
"fused_minimax_m3_qknorm_rope_kv_insert("
"Tensor! qkv, Tensor q_norm_weight, Tensor k_norm_weight, "
"Tensor cos_sin_cache, Tensor positions, int num_heads, "
"int num_kv_heads, int rotary_dim, float eps, "
"Tensor? index_q_norm_weight, Tensor? index_k_norm_weight, "
"int num_index_heads, "
"Tensor? slot_mapping, Tensor!? kv_cache, Tensor!? index_cache, "
"int block_size, Tensor!? q_out, Tensor!? index_q_out) -> ()");
// Apply repetition penalties to logits in-place.
ops.def(
"apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, "
@@ -364,9 +375,11 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
ops.def("mul_and_silu(Tensor! out, Tensor input) -> ()");
// SwiGLU activation with input clamping.
// alpha scales the sigmoid (gate * sigmoid(alpha * gate)); beta is added to
// the up half (up + beta). Defaults alpha=1.0, beta=0.0 give silu(gate)*up.
ops.def(
"silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit) "
"-> ()");
"silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit, "
"float alpha=1.0, float beta=0.0) -> ()");
// Activation function used in GeGLU with `none` approximation.
ops.def("gelu_and_mul(Tensor! out, Tensor input) -> ()");
@@ -571,6 +584,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
// Positional encoding kernels (shared CUDA/ROCm)
ops.impl("rotary_embedding", TORCH_BOX(&rotary_embedding));
ops.impl("fused_qk_norm_rope", TORCH_BOX(&fused_qk_norm_rope));
ops.impl("fused_minimax_m3_qknorm_rope_kv_insert",
TORCH_BOX(&fused_minimax_m3_qknorm_rope_kv_insert));
// Sampler kernels (shared CUDA/ROCm)
ops.impl("apply_repetition_penalties_",
+742
View File
@@ -0,0 +1,742 @@
// CUDA C++ q2k -> k2q CSR builder.
//
// Five-stage pipeline. q-ascending order within each CSR row is preserved
// by partitioning q across (CTA, warp_in_CTA) units; each unit owns a
// contiguous q-sub-range and reserves a contiguous slot range per row via
// a precomputed exclusive prefix scan.
//
// M: build_row_map -- round-robin packing of rows across batches
// H: histogram + tile_counts
// PR: row prefix -- single block per head, row_counts -> row_ptr
// PT: tile prefix -- multi-block, scan tile_counts along (c, w) axis
// S: scatter (sorted) -- per-warp slot range, q-sequential within warp
//
// Per-warp partitioning: each CTA has kWarps warps; warp w of CTA c owns
// q-range [c*q_per_cta + w*q_per_warp, c*q_per_cta + (w+1)*q_per_warp).
// tile_counts is shaped [G * kWarps, H, total_rows]; the "row" dimension
// of the prefix scan is the flattened (c * kWarps + w) index, scanned in
// lexicographic order so that warp-local slot ranges concatenate to the
// global q-sorted output.
#include <torch/all.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <algorithm>
#define CHECK_CUDA(x) TORCH_CHECK((x).is_cuda(), #x " must be CUDA")
#define CHECK_CONTIGUOUS(x) \
TORCH_CHECK((x).is_contiguous(), #x " must be contiguous")
#define CHECK_INT(x) \
TORCH_CHECK((x).scalar_type() == at::kInt, #x " must be int32")
#define CHECK_INPUT(x) \
CHECK_CUDA(x); \
CHECK_CONTIGUOUS(x); \
CHECK_INT(x)
namespace {
constexpr int kWarpSize = 32;
__device__ __forceinline__ void advance_batch_only(int const* __restrict__ cu_q,
int B, int q_abs, int& bi) {
while (bi < B && cu_q[bi + 1] <= q_abs) ++bi;
}
// Atomic increment of a 16-bit half within a 32-bit SMEM word; returns the
// OLD 16-bit value (slot). Per-warp count must stay < 32768 so the low
// half does not carry into the high half.
// base_int32 : int32 pointer; element i holds rows 2*i (low) and 2*i+1
// (high).
__device__ __forceinline__ int atomic_inc_int16_packed(int* base_int32,
int row) {
int idx = row >> 1;
int shift = (row & 1) << 4; // 0 or 16
int delta = 1 << shift;
int old = atomicAdd(&base_int32[idx], delta);
return (old >> shift) & 0xFFFF;
}
// Read 16-bit half from packed int32 storage.
__device__ __forceinline__ int read_int16_packed(int const* base_int32,
int row) {
int v = base_int32[row >> 1];
int shift = (row & 1) << 4;
return (v >> shift) & 0xFFFF;
}
// ---------------------------------------------------------------------------
// M: round-robin row map.
// ---------------------------------------------------------------------------
template <int kBlockK>
__global__ void k2q_build_row_map_kernel(int const* __restrict__ cu_k,
int* __restrict__ row_map,
int* __restrict__ row_coords, int B,
int max_kv_blocks) {
int level = blockIdx.x;
if (level >= max_kv_blocks) return;
if (threadIdx.x != 0) return;
int rows_before = 0;
for (int b = 0; b < B; ++b) {
int rb = (cu_k[b + 1] - cu_k[b] + kBlockK - 1) / kBlockK;
rows_before += (rb < level ? rb : level);
}
int active_before = 0;
for (int b = 0; b < B; ++b) {
int rb = (cu_k[b + 1] - cu_k[b] + kBlockK - 1) / kBlockK;
if (rb > level) {
int row_linear = rows_before + active_before;
row_map[(size_t)b * max_kv_blocks + level] = row_linear;
if (row_coords != nullptr) {
row_coords[(size_t)row_linear * 2] = b;
row_coords[(size_t)row_linear * 2 + 1] = level;
}
++active_before;
} else {
row_map[(size_t)b * max_kv_blocks + level] = -1;
}
}
}
// ---------------------------------------------------------------------------
// H: per-warp histogram + tile_counts.
// kWarps warps per CTA, each owns q-sub-range = q_per_cta / kWarps.
// SMEM hist[kWarps, total_rows] int32 (stored as packed int16 cursor:
// 2 entries per int32 word). Each warp counts to its own row.
// At end-of-CTA, write tile_counts[c*kWarps + w, h, r] = smem_hist[w, r]
// and atomicAdd(row_counts[h, r], sum over w of smem_hist[w, r]).
// ---------------------------------------------------------------------------
template <int kTopK, int kBlockK, int kWarps>
__global__ void k2q_hist_kernel(int const* __restrict__ q2k,
int const* __restrict__ cu_q,
int const* __restrict__ row_map,
int* __restrict__ row_counts,
int* __restrict__ tile_counts, int H, int B,
int S_Q, int total_rows, int max_kv_blocks,
int q_per_cta, int q_per_warp) {
constexpr int kThreads = kWarps * kWarpSize;
extern __shared__ int smem_hist_int[];
int* smem_hist = smem_hist_int;
int tid = threadIdx.x;
int warp_id = tid >> 5;
int lane = tid & 31;
int c = blockIdx.x;
int q_start_cta = c * q_per_cta;
int q_end_cta = min(q_start_cta + q_per_cta, S_Q);
int q_start_warp = min(q_start_cta + warp_id * q_per_warp, q_end_cta);
int q_end_warp = min(q_start_warp + q_per_warp, q_end_cta);
constexpr int kInt4PerToken = kTopK / 4;
int packed_per_warp = (total_rows + 1) >> 1;
int* my_hist = smem_hist + warp_id * packed_per_warp;
for (int h = 0; h < H; ++h) {
for (int i = lane; i < packed_per_warp; i += kWarpSize) my_hist[i] = 0;
__syncthreads();
if (q_start_warp < q_end_warp) {
int bi = 0;
int qi = q_start_warp + lane;
advance_batch_only(cu_q, B, qi, bi);
int4 const* head_topk4 =
reinterpret_cast<int4 const*>(q2k + (size_t)h * S_Q * kTopK);
for (; qi < q_end_warp; qi += kWarpSize) {
advance_batch_only(cu_q, B, qi, bi);
int const* my_row_map = row_map + (size_t)bi * max_kv_blocks;
int4 buf[kInt4PerToken];
#pragma unroll
for (int v = 0; v < kInt4PerToken; ++v) {
buf[v] = head_topk4[(size_t)qi * kInt4PerToken + v];
}
#pragma unroll
for (int t = 0; t < kTopK; ++t) {
int kvb_local = reinterpret_cast<int const*>(buf)[t];
if (kvb_local >= 0 && kvb_local < max_kv_blocks) {
int row = my_row_map[kvb_local];
if (row >= 0 && row < total_rows) {
atomic_inc_int16_packed(my_hist, row);
}
}
}
}
}
__syncthreads();
int* head_row_counts = row_counts + (size_t)h * total_rows;
// Each warp writes its own slice of tile_counts (full int32) by
// unpacking int16 entries from SMEM.
int* my_tile =
tile_counts + ((size_t)(c * kWarps + warp_id) * H + h) * total_rows;
for (int i = lane; i < total_rows; i += kWarpSize) {
my_tile[i] = read_int16_packed(my_hist, i);
}
__syncthreads();
// Sum across warps (int32 accumulator), atomicAdd to row_counts.
for (int i = tid; i < total_rows; i += kThreads) {
int sum = 0;
#pragma unroll
for (int w = 0; w < kWarps; ++w) {
sum += read_int16_packed(smem_hist + w * packed_per_warp, i);
}
if (sum > 0) atomicAdd(&head_row_counts[i], sum);
}
if (h + 1 < H) __syncthreads();
}
}
// ---------------------------------------------------------------------------
// PR: row prefix. One block per head.
// ---------------------------------------------------------------------------
template <int kThreads>
__global__ void k2q_row_prefix_kernel(int const* __restrict__ row_counts,
int* __restrict__ row_ptr,
int const* __restrict__ row_coords,
int* __restrict__ scheduler_metadata,
int* __restrict__ work_count,
int total_rows, int target_q_per_cta,
int work_capacity) {
int h = blockIdx.x;
int tid = threadIdx.x;
__shared__ int scan_buf[kThreads];
int const* head_counts = row_counts + (size_t)h * total_rows;
int* head_rowptr = row_ptr + (size_t)h * (total_rows + 1);
int chunk = (total_rows + kThreads - 1) / kThreads;
int lo = tid * chunk;
int hi = min(lo + chunk, total_rows);
int local_sum = 0;
for (int i = lo; i < hi; ++i) local_sum += head_counts[i];
scan_buf[tid] = local_sum;
__syncthreads();
for (int off = 1; off < kThreads; off <<= 1) {
int add = (tid >= off) ? scan_buf[tid - off] : 0;
__syncthreads();
scan_buf[tid] += add;
__syncthreads();
}
int running = scan_buf[tid] - local_sum;
for (int i = lo; i < hi; ++i) {
int row_count = head_counts[i];
running += row_count;
head_rowptr[i + 1] = running;
if (scheduler_metadata != nullptr && work_count != nullptr &&
row_count > 0) {
int num_chunks = (row_count + target_q_per_cta - 1) / target_q_per_cta;
int base = atomicAdd(work_count, num_chunks);
int batch_idx = row_coords[(size_t)i * 2];
int kv_block_idx = row_coords[(size_t)i * 2 + 1];
for (int c = 0; c < num_chunks; ++c) {
int work_idx = base + c;
if (work_idx < work_capacity) {
int q_begin = c * target_q_per_cta;
int q_count = min(target_q_per_cta, row_count - q_begin);
int* meta = scheduler_metadata + (size_t)work_idx * 6;
meta[0] = h;
meta[1] = i;
meta[2] = q_begin;
meta[3] = q_count;
meta[4] = batch_idx;
meta[5] = kv_block_idx;
}
}
}
}
}
// ---------------------------------------------------------------------------
// PT_smem: SMEM-staged tile prefix scan.
// Each block handles kRowsPerBlock rows for one head h. Cooperative load
// of tile_counts[*, h, base_r..base_r+M) into SMEM (better coalescing
// than per-warp uncoalesced stride reads), then per-warp scan in SMEM,
// then cooperative store back. Fuses row_ptr into the base.
// ---------------------------------------------------------------------------
template <int kThreads, int kRowsPerBlock>
__global__ void k2q_tile_prefix_smem_kernel(int* __restrict__ tile_counts,
int const* __restrict__ row_ptr,
int H, int total_rows,
int G_total) {
static_assert(kRowsPerBlock > 0, "kRowsPerBlock must be positive");
extern __shared__ int smem_tprefix[];
// smem layout: smem[r_off][g] for r_off in [0, M), g in [0, G_total).
int tid = threadIdx.x;
int lane = tid & 31;
int warp_id = tid >> 5;
// Grid: H * blocks_per_h. Each block stays within a single head h
// and processes kRowsPerBlock contiguous rows starting at b_in_h *
// kRowsPerBlock. (Earlier flat-grid mapping `h = block_job /
// total_rows; base_r = block_job - h*total_rows` skipped rows when
// total_rows was not a multiple of kRowsPerBlock and H > 1, because
// the last partial block of head h-1 left blocks of head h starting
// at a non-zero row offset.)
int blocks_per_h = (total_rows + kRowsPerBlock - 1) / kRowsPerBlock;
int h = blockIdx.x / blocks_per_h;
int b_in_h = blockIdx.x - h * blocks_per_h;
if (h >= H) return;
int base_r = b_in_h * kRowsPerBlock;
if (base_r >= total_rows) return;
int actual_M = min(kRowsPerBlock, total_rows - base_r);
size_t stride_g = (size_t)H * total_rows;
int* base_ptr = tile_counts + (size_t)h * total_rows + base_r;
int total_elems = G_total * actual_M;
// Cooperative load. Pattern: thread tid -> (r_off=tid%M, g=tid/M),
// then strided. 32 lanes hit M r's × (32/M) g's, giving 32/M cache
// lines per warp (vs 32 in the naive stride-along-g pattern).
for (int i = tid; i < total_elems; i += kThreads) {
int r_off = i % actual_M;
int g = i / actual_M;
smem_tprefix[r_off * G_total + g] = base_ptr[g * stride_g + r_off];
}
__syncthreads();
// Per-warp scan: warp w scans row (base_r + w) if w < actual_M.
if (warp_id < actual_M) {
int abs_r = base_r + warp_id;
int rp = row_ptr[(size_t)h * (total_rows + 1) + abs_r];
int* my_smem = smem_tprefix + warp_id * G_total;
int running = rp;
for (int g0 = 0; g0 < G_total; g0 += kWarpSize) {
int g = g0 + lane;
int v = (g < G_total) ? my_smem[g] : 0;
int x = v;
#pragma unroll
for (int off = 1; off < kWarpSize; off <<= 1) {
int nbr = __shfl_up_sync(0xFFFFFFFF, x, off);
if (lane >= off) x += nbr;
}
int excl = running + x - v;
if (g < G_total) my_smem[g] = excl;
int chunk_sum = __shfl_sync(0xFFFFFFFF, x, 31);
running += chunk_sum;
}
}
__syncthreads();
// Cooperative store back.
for (int i = tid; i < total_elems; i += kThreads) {
int r_off = i % actual_M;
int g = i / actual_M;
base_ptr[g * stride_g + r_off] = smem_tprefix[r_off * G_total + g];
}
}
// ---------------------------------------------------------------------------
// S: scatter. kWarps warps per CTA, each owns q-sub-range. Per-warp SMEM
// cursor and per-warp tile_offset slot range. Within a warp, q's are
// processed sequentially; lanes 0..kTopK-1 handle the topK slots in
// lockstep. Across distinct q's in the same warp, the lockstep ordering
// guarantees q-monotonic atomicAdd on smem_cursor[r].
// ---------------------------------------------------------------------------
// kQPerIter * kTopK lanes are active per warp iter; remaining lanes idle.
// For kTopK=16, kQPerIter=2 uses all 32 lanes; for kTopK=8, kQPerIter=4.
// CORRECTNESS NOTE: relies on lane-ordered SMEM atomicAdd return values
// within a single warp instruction (verified on B200; tests pass).
//
// SMEM cursor stored as packed int16 (two cursors per int32). Per-warp
// row count must stay < 32768 (~q_per_warp * kTopK at max sink), which
// holds for all task.md sizes up to 1024K.
template <int kTopK, int kBlockK, int kWarps>
__global__ void k2q_scatter_kernel(
int const* __restrict__ q2k, int const* __restrict__ cu_q,
int const* __restrict__ row_map, int const* __restrict__ abs_base,
int* __restrict__ q_idx, int* __restrict__ qsplit_idx,
int* __restrict__ split_counts, int H, int B, int S_Q, int total_rows,
int max_kv_blocks, int q_per_cta, int q_per_warp, int max_seqlen_q) {
constexpr int kQPerIter = kWarpSize / kTopK > 0 ? kWarpSize / kTopK : 1;
extern __shared__ int smem_cursor_int[];
int* smem_cursor = smem_cursor_int;
int tid = threadIdx.x;
int warp_id = tid >> 5;
int lane = tid & 31;
int c = blockIdx.x;
int q_start_cta = c * q_per_cta;
int q_end_cta = min(q_start_cta + q_per_cta, S_Q);
int q_start_warp = min(q_start_cta + warp_id * q_per_warp, q_end_cta);
int q_end_warp = min(q_start_warp + q_per_warp, q_end_cta);
int q_in_iter = lane / kTopK;
int slot_in_q = lane % kTopK;
bool lane_active = (lane < kQPerIter * kTopK);
// Per-warp packed cursor: total_rows int16 entries -> ceil(total_rows/2)
// int32.
int packed_per_warp = (total_rows + 1) >> 1;
int* my_cursor = smem_cursor + warp_id * packed_per_warp;
for (int h = 0; h < H; ++h) {
for (int i = lane; i < packed_per_warp; i += kWarpSize) my_cursor[i] = 0;
__syncwarp();
if (q_start_warp < q_end_warp) {
int bi = 0;
advance_batch_only(cu_q, B, q_start_warp, bi);
int const* head_q2k = q2k + (size_t)h * S_Q * kTopK;
int const* my_abs_base =
abs_base + ((size_t)(c * kWarps + warp_id) * H + h) * total_rows;
int* head_qidx = q_idx + (size_t)h * S_Q * kTopK;
// (Hot-row register cache experiment showed no measurable
// benefit; relying on L1 to keep row 0 / row total_rows-1
// hot since they're hit every iteration in sink workloads.)
constexpr int kUnroll = 16;
int qi_base = q_start_warp;
for (; qi_base + kUnroll * kQPerIter <= q_end_warp;
qi_base += kUnroll * kQPerIter) {
int kvb[kUnroll];
int qloc[kUnroll];
int batch[kUnroll];
int const* rmap[kUnroll];
#pragma unroll
for (int u = 0; u < kUnroll; ++u) {
int qi_u = qi_base + u * kQPerIter + q_in_iter;
kvb[u] = -1;
qloc[u] = 0;
batch[u] = 0;
if (lane_active) {
advance_batch_only(cu_q, B, qi_u, bi);
qloc[u] = qi_u - cu_q[bi];
batch[u] = bi;
kvb[u] = head_q2k[(size_t)qi_u * kTopK + slot_in_q];
}
rmap[u] = row_map + (size_t)bi * max_kv_blocks;
}
int row[kUnroll];
#pragma unroll
for (int u = 0; u < kUnroll; ++u) {
row[u] = -1;
if (lane_active && kvb[u] >= 0 && kvb[u] < max_kv_blocks)
row[u] = rmap[u][kvb[u]];
}
// Pre-issue all kUnroll abs_base loads in parallel before
// the atomic chain so memory pipeline runs concurrently
// with SMEM atomic-adds.
int abs_v[kUnroll];
#pragma unroll
for (int u = 0; u < kUnroll; ++u) {
abs_v[u] =
(row[u] >= 0 && row[u] < total_rows) ? my_abs_base[row[u]] : 0;
}
#pragma unroll
for (int u = 0; u < kUnroll; ++u) {
int r = row[u];
bool valid_edge = r >= 0 && r < total_rows;
unsigned int valid_mask = __ballot_sync(0xFFFFFFFFu, valid_edge);
unsigned int group_mask =
(kTopK == 32) ? 0xFFFFFFFFu
: (((1u << kTopK) - 1u) << (q_in_iter * kTopK));
unsigned int lower_lane_mask = lane == 0 ? 0u : ((1u << lane) - 1u);
int split_slot = __popc(valid_mask & group_mask & lower_lane_mask);
int valid_count = __popc(valid_mask & group_mask);
if (split_counts != nullptr && slot_in_q == 0) {
split_counts[((size_t)batch[u] * max_seqlen_q + qloc[u]) * H + h] =
valid_count;
}
if (valid_edge) {
int slot = atomic_inc_int16_packed(my_cursor, r);
int out_pos = abs_v[u] + slot;
head_qidx[out_pos] = qloc[u];
if (qsplit_idx != nullptr) {
qsplit_idx[(size_t)h * S_Q * kTopK + out_pos] =
qloc[u] | ((split_slot & 0xFF) << 24);
}
}
}
}
// Tail: 1-3 iters left.
for (; qi_base < q_end_warp; qi_base += kQPerIter) {
int my_qi = qi_base + q_in_iter;
bool valid_q = (my_qi < q_end_warp) && lane_active;
int kvb_local = -1;
int q_local = 0;
int batch_local = 0;
if (valid_q) {
advance_batch_only(cu_q, B, my_qi, bi);
batch_local = bi;
q_local = my_qi - cu_q[bi];
kvb_local = head_q2k[(size_t)my_qi * kTopK + slot_in_q];
}
int const* my_row_map = row_map + (size_t)bi * max_kv_blocks;
int row = -1;
if (valid_q && kvb_local >= 0 && kvb_local < max_kv_blocks) {
row = my_row_map[kvb_local];
}
bool valid_edge = row >= 0 && row < total_rows;
unsigned int valid_mask = __ballot_sync(0xFFFFFFFFu, valid_edge);
unsigned int group_mask =
(kTopK == 32) ? 0xFFFFFFFFu
: (((1u << kTopK) - 1u) << (q_in_iter * kTopK));
unsigned int lower_lane_mask = lane == 0 ? 0u : ((1u << lane) - 1u);
int split_slot = __popc(valid_mask & group_mask & lower_lane_mask);
int valid_count = __popc(valid_mask & group_mask);
if (split_counts != nullptr && valid_q && slot_in_q == 0) {
split_counts[((size_t)batch_local * max_seqlen_q + q_local) * H + h] =
valid_count;
}
if (valid_edge) {
int slot = atomic_inc_int16_packed(my_cursor, row);
int out_pos = my_abs_base[row] + slot;
head_qidx[out_pos] = q_local;
if (qsplit_idx != nullptr) {
qsplit_idx[(size_t)h * S_Q * kTopK + out_pos] =
q_local | ((split_slot & 0xFF) << 24);
}
}
}
}
if (h + 1 < H) __syncthreads();
}
}
} // anonymous namespace
// ===========================================================================
// Host orchestration
// ===========================================================================
template <int kTopK, int kBlockK>
static void launch_pipeline(torch::Tensor q2k, torch::Tensor cu_q,
torch::Tensor cu_k, torch::Tensor row_ptr,
torch::Tensor q_idx, int total_rows,
int max_kv_blocks,
torch::Tensor scheduler_metadata = torch::Tensor(),
torch::Tensor work_count = torch::Tensor(),
torch::Tensor qsplit_idx = torch::Tensor(),
torch::Tensor split_counts = torch::Tensor(),
int target_q_per_cta = 1, int work_capacity = 0,
int max_seqlen_q = 0) {
int H = (int)q2k.size(0);
int S_Q = (int)q2k.size(1);
int topK = (int)q2k.size(2);
TORCH_CHECK(topK == kTopK, "topK runtime != template kTopK");
int B = (int)cu_q.size(0) - 1;
auto device = q2k.device();
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
AT_CUDA_CHECK(cudaMemsetAsync(row_ptr.data_ptr<int>(), 0,
(size_t)H * (total_rows + 1) * sizeof(int),
stream));
AT_CUDA_CHECK(cudaMemsetAsync(q_idx.data_ptr<int>(), 0xFF,
(size_t)H * S_Q * kTopK * sizeof(int), stream));
auto opts = torch::TensorOptions().dtype(torch::kInt32).device(device);
auto row_counts = torch::zeros({H, total_rows}, opts);
auto row_map = torch::empty({B, max_kv_blocks}, opts);
bool emit_schedule = scheduler_metadata.defined();
auto row_coords =
emit_schedule ? torch::empty({total_rows, 2}, opts) : torch::Tensor();
int* scheduler_metadata_ptr =
emit_schedule ? scheduler_metadata.data_ptr<int>() : nullptr;
int* work_count_ptr = emit_schedule ? work_count.data_ptr<int>() : nullptr;
int* qsplit_idx_ptr = emit_schedule ? qsplit_idx.data_ptr<int>() : nullptr;
int* split_counts_ptr =
emit_schedule ? split_counts.data_ptr<int>() : nullptr;
int* row_coords_ptr = emit_schedule ? row_coords.data_ptr<int>() : nullptr;
if (emit_schedule) {
AT_CUDA_CHECK(cudaMemsetAsync(work_count_ptr, 0, sizeof(int), stream));
AT_CUDA_CHECK(cudaMemsetAsync(scheduler_metadata_ptr, 0,
(size_t)work_capacity * 6 * sizeof(int),
stream));
}
int dev = q2k.get_device();
int num_sms = 0;
AT_CUDA_CHECK(
cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, dev));
// -- Pick kWarps per CTA based on SMEM budget for cursor/hist ---------
// SMEM per CTA = kWarps * total_rows * sizeof(int) (for both H and S).
// Want at least 2 CTAs/SM for memory parallelism. SM100 SMEM = 228KB.
// Pick the largest kWarps that fits two CTAs/SM, capped at 4.
// SMEM cursor packed as int16 (2 entries per int32 word):
int per_warp_smem = ((total_rows + 1) >> 1) * (int)sizeof(int);
int kWarps_pick = 4;
while (kWarps_pick > 1 && (kWarps_pick * per_warp_smem) * 2 > 228 * 1024) {
kWarps_pick >>= 1;
}
if (kWarps_pick < 1) kWarps_pick = 1;
// -- Pick G (CTAs) ----------------------------------------------------
// For each (kWarps, per_warp_smem) pair, the SMEM-bound occupancy is
// 228KB / (kWarps*per_warp_smem) CTAs/SM. We size G as
// num_sms * occupancy so a single resident wave covers all CTAs and
// the memory pipeline runs at peak.
int per_cta_smem_bytes = kWarps_pick * per_warp_smem;
int max_ctas_per_sm =
std::max(1, (228 * 1024) / std::max(1, per_cta_smem_bytes));
if (max_ctas_per_sm > 8) max_ctas_per_sm = 8;
constexpr int kMinQPerCta = 256;
// Cap target_g at num_sms * 3 — empirically this balances
// per-CTA work-size against parallelism. Higher caps regress
// mid-size cases due to row_counts atomicAdd contention and
// smaller q_per_cta. SMEM-bound configurations naturally cap
// lower if max_ctas_per_sm < 3.
int target_g = num_sms * std::min(max_ctas_per_sm, 3);
int max_g_for_q = (S_Q + kMinQPerCta - 1) / kMinQPerCta;
int G = std::min({target_g, max_g_for_q, S_Q});
if (G < 1) G = 1;
int q_per_cta = (S_Q + G - 1) / G;
G = (S_Q + q_per_cta - 1) / q_per_cta;
int q_per_warp = (q_per_cta + kWarps_pick - 1) / kWarps_pick;
int G_total = G * kWarps_pick;
auto tile_counts = torch::empty({G_total, H, total_rows}, opts);
// -- Compile-time switch on kWarps for the templated kernels ---------
auto rmap_fn = k2q_build_row_map_kernel<kBlockK>;
auto rprefix_fn = k2q_row_prefix_kernel<1024>;
constexpr int kPtRowsPerBlock = 8;
constexpr int kPtThreads = 256;
auto tprefix_smem_fn =
k2q_tile_prefix_smem_kernel<kPtThreads, kPtRowsPerBlock>;
if (max_kv_blocks > 0) {
rmap_fn<<<max_kv_blocks, 32, 0, stream>>>(cu_k.data_ptr<int>(),
row_map.data_ptr<int>(),
row_coords_ptr, B, max_kv_blocks);
}
auto launch_hist_scatter = [&](auto kWarps_const) {
constexpr int W = decltype(kWarps_const)::value;
size_t smem_bytes = (size_t)W * per_warp_smem;
auto hist_fn = k2q_hist_kernel<kTopK, kBlockK, W>;
auto scat_fn = k2q_scatter_kernel<kTopK, kBlockK, W>;
AT_CUDA_CHECK(cudaFuncSetAttribute(
hist_fn, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem_bytes));
AT_CUDA_CHECK(cudaFuncSetAttribute(
scat_fn, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem_bytes));
hist_fn<<<G, W * kWarpSize, smem_bytes, stream>>>(
q2k.data_ptr<int>(), cu_q.data_ptr<int>(), row_map.data_ptr<int>(),
row_counts.data_ptr<int>(), tile_counts.data_ptr<int>(), H, B, S_Q,
total_rows, max_kv_blocks, q_per_cta, q_per_warp);
rprefix_fn<<<H, 1024, 0, stream>>>(
row_counts.data_ptr<int>(), row_ptr.data_ptr<int>(),
emit_schedule ? row_coords.data_ptr<int>() : nullptr,
scheduler_metadata_ptr, work_count_ptr, total_rows, target_q_per_cta,
work_capacity);
// Grid is H * blocks_per_h so each block stays within a single
// head; flat (H*total_rows) grid would skip rows when total_rows
// is not a multiple of kPtRowsPerBlock.
int blocks_per_h = (total_rows + kPtRowsPerBlock - 1) / kPtRowsPerBlock;
int pt_grid = H * blocks_per_h;
if (pt_grid < 1) pt_grid = 1;
size_t pt_smem = (size_t)kPtRowsPerBlock * G_total * sizeof(int);
AT_CUDA_CHECK(cudaFuncSetAttribute(
tprefix_smem_fn, cudaFuncAttributeMaxDynamicSharedMemorySize,
(int)pt_smem));
tprefix_smem_fn<<<pt_grid, kPtThreads, pt_smem, stream>>>(
tile_counts.data_ptr<int>(), row_ptr.data_ptr<int>(), H, total_rows,
G_total);
scat_fn<<<G, W * kWarpSize, smem_bytes, stream>>>(
q2k.data_ptr<int>(), cu_q.data_ptr<int>(), row_map.data_ptr<int>(),
tile_counts.data_ptr<int>(), q_idx.data_ptr<int>(), qsplit_idx_ptr,
split_counts_ptr, H, B, S_Q, total_rows, max_kv_blocks, q_per_cta,
q_per_warp, max_seqlen_q);
};
if (kWarps_pick == 4) {
launch_hist_scatter(std::integral_constant<int, 4>{});
} else if (kWarps_pick == 2) {
launch_hist_scatter(std::integral_constant<int, 2>{});
} else {
launch_hist_scatter(std::integral_constant<int, 1>{});
}
}
void run_minimax_m3_build_k2q_csr_with_schedule(
torch::Tensor q2k, torch::Tensor cu_q, torch::Tensor cu_k,
torch::Tensor row_ptr, torch::Tensor q_idx,
torch::Tensor scheduler_metadata, torch::Tensor work_count,
torch::Tensor qsplit_idx, torch::Tensor split_counts, int64_t topk,
int64_t blk_kv, int64_t total_rows, int64_t max_kv_blocks,
int64_t target_q_per_cta, int64_t work_capacity, int64_t max_seqlen_q) {
CHECK_INPUT(q2k);
CHECK_INPUT(cu_q);
CHECK_INPUT(cu_k);
CHECK_INPUT(row_ptr);
CHECK_INPUT(q_idx);
CHECK_INPUT(scheduler_metadata);
CHECK_INPUT(work_count);
CHECK_INPUT(qsplit_idx);
CHECK_INPUT(split_counts);
TORCH_CHECK(blk_kv == 128, "build_k2q_csr only supports blk_kv == 128");
int H = (int)q2k.size(0);
int S_Q = (int)q2k.size(1);
int tr = (int)total_rows;
int mkv = (int)max_kv_blocks;
int target = (int)target_q_per_cta;
int capacity = (int)work_capacity;
int max_sq = (int)max_seqlen_q;
TORCH_CHECK(tr >= 0 && mkv >= 0 && target > 0 && capacity > 0 && max_sq >= 0,
"invalid schedule sizing arguments");
TORCH_CHECK(row_ptr.size(0) == H && row_ptr.size(1) == tr + 1,
"row_ptr shape mismatch");
TORCH_CHECK(q_idx.size(0) == H && q_idx.size(1) == (int64_t)S_Q * (int)topk,
"q_idx shape mismatch");
TORCH_CHECK(qsplit_idx.sizes() == q_idx.sizes(), "qsplit_idx shape mismatch");
TORCH_CHECK(
scheduler_metadata.size(0) == capacity && scheduler_metadata.size(1) == 6,
"scheduler_metadata shape mismatch");
TORCH_CHECK(work_count.numel() == 1,
"work_count must have one int32 element");
TORCH_CHECK(split_counts.dim() == 3 &&
split_counts.size(0) == cu_q.size(0) - 1 &&
split_counts.size(1) == max_sq && split_counts.size(2) == H,
"split_counts shape mismatch");
if (S_Q == 0 || tr == 0 || H == 0 || mkv == 0) {
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
AT_CUDA_CHECK(cudaMemsetAsync(row_ptr.data_ptr<int>(), 0,
(size_t)H * (tr + 1) * sizeof(int), stream));
AT_CUDA_CHECK(cudaMemsetAsync(q_idx.data_ptr<int>(), 0xFF,
(size_t)H * S_Q * (int)topk * sizeof(int),
stream));
AT_CUDA_CHECK(
cudaMemsetAsync(work_count.data_ptr<int>(), 0, sizeof(int), stream));
if (split_counts.numel() > 0) {
AT_CUDA_CHECK(cudaMemsetAsync(split_counts.data_ptr<int>(), 0,
(size_t)split_counts.numel() * sizeof(int),
stream));
}
return;
}
if (topk == 16) {
launch_pipeline<16, 128>(q2k, cu_q, cu_k, row_ptr, q_idx, tr, mkv,
scheduler_metadata, work_count, qsplit_idx,
split_counts, target, capacity, max_sq);
} else if (topk == 8) {
launch_pipeline<8, 128>(q2k, cu_q, cu_k, row_ptr, q_idx, tr, mkv,
scheduler_metadata, work_count, qsplit_idx,
split_counts, target, capacity, max_sq);
} else if (topk == 32) {
launch_pipeline<32, 128>(q2k, cu_q, cu_k, row_ptr, q_idx, tr, mkv,
scheduler_metadata, work_count, qsplit_idx,
split_counts, target, capacity, max_sq);
} else if (topk == 4) {
launch_pipeline<4, 128>(q2k, cu_q, cu_k, row_ptr, q_idx, tr, mkv,
scheduler_metadata, work_count, qsplit_idx,
split_counts, target, capacity, max_sq);
} else {
TORCH_CHECK(false, "unsupported topK ", topk,
" (expected 4, 8, 16, or 32)");
}
}
+10 -1
View File
@@ -62,7 +62,8 @@ void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
void silu_and_mul(torch::Tensor& out, torch::Tensor& input);
void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit);
void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit,
double alpha = 1.0, double beta = 0.0);
void silu_and_mul_quant(torch::Tensor& out, torch::Tensor& input,
torch::Tensor& scale);
@@ -146,4 +147,12 @@ std::tuple<torch::Tensor, torch::Tensor> minimax_allreduce_rms_qk(
torch::Tensor const& norm_weight_k, torch::Tensor workspace,
int64_t const q_size, int64_t const kv_size, int64_t const rank,
int64_t const nranks, double const eps);
void run_minimax_m3_build_k2q_csr_with_schedule(
torch::Tensor q2k, torch::Tensor cu_q, torch::Tensor cu_k,
torch::Tensor row_ptr, torch::Tensor q_idx,
torch::Tensor scheduler_metadata, torch::Tensor work_count,
torch::Tensor qsplit_idx, torch::Tensor split_counts, int64_t topk,
int64_t blk_kv, int64_t total_rows, int64_t max_kv_blocks,
int64_t target_q_per_cta, int64_t work_capacity, int64_t max_seqlen_q);
#endif
+21
View File
@@ -187,6 +187,27 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
"float eps) -> (Tensor, Tensor)");
ops.impl("minimax_allreduce_rms_qk", torch::kCUDA, &minimax_allreduce_rms_qk);
ops.def(
"minimax_m3_build_k2q_csr_with_schedule("
"Tensor q2k,"
"Tensor cu_q,"
"Tensor cu_k,"
"Tensor! row_ptr,"
"Tensor! q_idx,"
"Tensor! scheduler_metadata,"
"Tensor! work_count,"
"Tensor! qsplit_idx,"
"Tensor! split_counts,"
"int topk,"
"int blk_kv,"
"int total_rows,"
"int max_kv_blocks,"
"int target_q_per_cta,"
"int work_capacity,"
"int max_seqlen_q) -> ()");
ops.impl("minimax_m3_build_k2q_csr_with_schedule", torch::kCUDA,
&run_minimax_m3_build_k2q_csr_with_schedule);
// conditionally compiled so impl in source file
#endif
}
+1 -1
View File
@@ -757,7 +757,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
# Install FlashInfer JIT cache (requires CUDA-version-specific index URL)
# https://docs.flashinfer.ai/installation.html
# From versions.json: .flashinfer.version
ARG FLASHINFER_VERSION=0.6.11.post2
ARG FLASHINFER_VERSION=0.6.12
RUN --mount=type=cache,target=/opt/uv/cache \
uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
--extra-index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
+2 -2
View File
@@ -256,13 +256,13 @@ RUN pip install setuptools==75.6.0 packaging==23.2 ninja==1.11.1.3 build==1.2.2.
# build flashinfer for torch nightly from source around 10 mins
# release version: v0.6.11.post2
# release version: v0.6.12
# todo(elainewy): cache flashinfer build result for faster build
ENV CCACHE_DIR=/root/.cache/ccache
RUN --mount=type=cache,target=/root/.cache/ccache \
--mount=type=cache,target=/root/.cache/uv \
echo "git clone flashinfer..." \
&& git clone --depth 1 --branch v0.6.11.post2 --recursive https://github.com/flashinfer-ai/flashinfer.git \
&& git clone --depth 1 --branch v0.6.12 --recursive https://github.com/flashinfer-ai/flashinfer.git \
&& cd flashinfer \
&& git submodule update --init --recursive \
&& echo "finish git clone flashinfer..." \
+1 -1
View File
@@ -9,7 +9,7 @@ ARG PYTORCH_AUDIO_BRANCH="v2.9.0"
ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git"
ARG FA_BRANCH="0e60e394"
ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
ARG AITER_BRANCH="v0.1.13.post1"
ARG AITER_BRANCH="v0.1.13"
ARG AITER_REPO="https://github.com/ROCm/aiter.git"
ARG MORI_BRANCH="v1.1.0"
ARG MORI_REPO="https://github.com/ROCm/mori.git"
+1 -1
View File
@@ -68,7 +68,7 @@
"default": "true"
},
"FLASHINFER_VERSION": {
"default": "0.6.11.post2"
"default": "0.6.12"
},
"GDRCOPY_CUDA_VERSION": {
"default": "12.8"
+2 -8
View File
@@ -246,12 +246,6 @@ Every image listed in "image_files" is added to the request in the listed order
The "image" shorthand accepts the same values as "image_files". The "image_url" field accepts either an OpenAI-style object with a "url" field or a URL string.
By default, image references are sent to the serving endpoint as provided, with local image paths converted to `file://` URLs.
If the benchmark client should load local and HTTP(S) images before sending requests, pass `--custom-ensure-client-side-data` to encode them as base64 data URLs on the client side.
Existing `data:image/...` URLs are already self-contained and are kept unchanged.
```bash
# need a model with vision capability here
vllm serve Qwen/Qwen2-VL-7B-Instruct
@@ -259,13 +253,13 @@ vllm serve Qwen/Qwen2-VL-7B-Instruct
```bash
# run benchmarking script
vllm bench serve --save-result --save-detailed \
vllm bench serve--save-result --save-detailed \
--backend openai-chat \
--model Qwen/Qwen2-VL-7B-Instruct \
--endpoint /v1/chat/completions \
--dataset-name custom_image \
--dataset-path <path-to-your-image-data-jsonl> \
--custom-ensure-client-side-data
--allowed-local-media-path /path/to/image/folder
```
Note that we need to use the `openai-chat` backend and `/v1/chat/completions` endpoint for multimodal inputs.
+14 -2
View File
@@ -170,8 +170,8 @@ Priority is **1 = highest** (tried first).
| Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. |
| ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ |
| `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | ❌ | All | N/A |
| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256, 512 | ❌ | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64 | 64, 128, 256, 512 | ✅ | ❌ | ❌ | ✅ | Decoder | 10.x |
| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | ❌ | ❌ | ✅ | Decoder | 10.x |
| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 |
| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x |
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
@@ -187,6 +187,18 @@ Priority is **1 = highest** (tried first).
>
> **\*** Specify the FlashAttention version via `--attention-config.flash_attn_version=2`, `3`, or `4`. Default is FA4 on SM100+ (Blackwell), FA3 on SM90 (Hopper), FA2 otherwise.
## MiniMax M3 Sparse Attention Backends
Block-sparse GQA backend used by MiniMax M3 sparse ("lightning indexer")
layers. It is wired in directly by the model and is not part of the
automatic priority lists above. A lightning indexer scores KV blocks, the
top-k blocks (plus fixed init/local blocks) are selected, and attention
attends only to those blocks; index keys live in a separate side cache.
| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. |
| ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ |
| `MINIMAX_M3_SPARSE` | bf16, fp16 | `bfloat16` | 128 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
## MLA (Multi-head Latent Attention) Backends
MLA uses separate backends for prefill and decode phases.
+1 -1
View File
@@ -778,7 +778,7 @@ Then, you can use the OpenAI client as follows:
base_url=openai_api_base,
)
video_url = "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
video_url = "http://commondatastorage.googleapis.com/gtv-videos-bucket/sample/ForBiggerFun.mp4"
## Use video url in the payload
chat_completion_from_url = client.chat.completions.create(
-1
View File
@@ -437,7 +437,6 @@ th {
| `LongcatFlashForCausalLM` | LongCat-Flash | `meituan-longcat/LongCat-Flash-Chat`, `meituan-longcat/LongCat-Flash-Chat-FP8` | ✅︎ | ✅︎ |
| `MambaForCausalLM` | Mamba | `state-spaces/mamba-130m-hf`, `state-spaces/mamba-790m-hf`, `state-spaces/mamba-2.8b-hf`, etc. | | ✅︎ |
| `Mamba2ForCausalLM` | Mamba2 | `mistralai/Mamba-Codestral-7B-v0.1`, etc. | | ✅︎ |
| `MellumForCausalLM` | Mellum 2 | `JetBrains/Mellum2-12B-A2.5B-Base`, etc. | | ✅︎ |
| `MiMoForCausalLM` | MiMo | `XiaomiMiMo/MiMo-7B-RL`, etc. | ✅︎ | ✅︎ |
| `MiMoV2FlashForCausalLM` | MiMoV2Flash | `XiaomiMiMo/MiMo-V2-Flash`, etc. | | ✅︎ |
| `MiMoV2ForCausalLM` | MiMoV2Pro | `XiaomiMiMo/MiMo-V2.5-Pro`, etc. | | ✅︎ |
+2 -32
View File
@@ -100,44 +100,14 @@ For further details on renderer APIs, please refer to [this page](renderer.md).
- `/version` - Version information
- `/load` - Server load metrics
## Server in development mode
When using the flag VLLM_SERVER_DEV_MODE=1, you enable development endpoints.
**SECURITY WARNING: These endpoints should NOT be used in production!**
### Cache Management APIs
- `/reset_prefix_cache` - Reset prefix cache (can disrupt service)
- `/reset_mm_cache` - Reset multimodal cache (can disrupt service)
- `/reset_encoder_cache` - Reset encoder cache (can disrupt service)
### Weight Transfer APIs (RL Training)
For further details on Weight Transfer, please refer to [this page](../../training/weight_transfer/README.md).
- `/pause` - Pause generation (causes denial of service)
- `/resume` - Resume generation
- `/is_paused` - Check if generation is paused
- `/init_weight_transfer_engine` - Initialize weight transfer engine for RLHF
- `/update_weights` - Update model weights (can alter model behavior)
- `/get_world_size` - Get distributed world size
### Collective RPC
- `/collective_rpc` - Execute arbitrary RPC methods on the engine (extremely dangerous)
### Server info
- `/server_info` - Get detailed server configuration
### Sleep Mode APIs
## Sleep Mode APIs
For further details on sleep mode, please refer to [this page](../../features/sleep_mode.md).
- `/sleep` - Put engine to sleep (causes denial of service)
- `/wake_up` - Wake engine from sleep
- `/is_sleeping` - Check if engine is sleeping
- `/collective_rpc` - Execute arbitrary RPC methods on the engine (extremely dangerous)
## Chat Template
+3 -18
View File
@@ -84,10 +84,7 @@ Both the trainer (`NCCLTrainerSendWeightsArgs`) and inference side (`NCCLWeightT
## Receiving Weights (Inference Side)
The inference side triggers weight reception using the four-phase protocol:
`init_weight_transfer_engine`, `start_weight_update`, `update_weights`,
`finish_weight_update`. The init phase is shown [above](#initialization). The
remaining three steps are:
The inference side triggers weight reception using the four-phase protocol`init_weight_transfer_engine`, `start_weight_update`, `update_weights`, `finish_weight_update`. The init phase is shown [above](#initialization). The remaining three steps are:
```python
from vllm.distributed.weight_transfer.base import WeightTransferUpdateRequest
@@ -111,24 +108,12 @@ llm.update_weights(
llm.finish_weight_update()
```
The `names`, `dtype_names`, and `shapes` lists describe each parameter. These
must match the order in which the trainer iterates over its parameters.
The `names`, `dtype_names`, and `shapes` lists describe each parameter. These must match the order in which the trainer iterates over its parameters.
`start_weight_update` must be called before `update_weights`, and
`finish_weight_update` must be called after all weight chunks have been
transferred. The `is_checkpoint_format` flag controls whether layerwise reload
processing is applied (`True` for checkpoint-format weights, `False` for
pre-processed kernel-format weights).
Sparse NCCL patches still use `update_kind="sparse_flat"` inside
`update_info`, but they should be wrapped in
`start_weight_update(is_checkpoint_format=False)` because sparse patches apply
directly to runtime/kernel-format parameters. The current sparse MVP requires
`TP=1` and `PP=1`.
`start_weight_update` must be called before `update_weights`, and `finish_weight_update` must be called after all weight chunks have been transferred. The `is_checkpoint_format` flag controls whether layerwise reload processing is applied (`True` for checkpoint-format weights, `False` for pre-processed kernel-format weights).
## Examples
- [RLHF with NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_nccl.py) - Trainer on one GPU, 2x tensor-parallel vLLM engine on two others, with packed NCCL weight broadcast
- [RLHF with sparse NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_sparse_nccl.py) - Dense-vs-sparse equivalence demo with a real model on a 2-GPU trainer/inference setup; sparse patches use `start_weight_update(is_checkpoint_format=False)` and currently require `TP=1` and `PP=1`
- [RLHF with async weight syncing (offline, Ray)](../../../examples/rl/rlhf_async_new_apis.py) - Async generation with mid-flight pause, weight sync, resume, and validation against a fresh model
- [RLHF with NCCL weight syncing (online serving, HTTP)](../../../examples/rl/rlhf_http_nccl.py) - Weight transfer with a running vLLM HTTP server using HTTP control plane and NCCL data plane
@@ -203,7 +203,7 @@ def run_multi_image(model: str, max_completion_tokens: int) -> None:
# Video input inference
def run_video(model: str, max_completion_tokens: int) -> None:
video_url = "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
video_url = "http://commondatastorage.googleapis.com/gtv-videos-bucket/sample/ForBiggerFun.mp4"
video_base64 = encode_base64_content_from_url(video_url)
## Use video url in the payload
-526
View File
@@ -1,526 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Demonstrates dense-vs-sparse NCCL weight syncing with a real model.
This example mirrors the validation story used for the sparse NCCL MVP:
both the dense update path and the sparse patch path start from the same real
checkpoint and apply the same deterministic trainer-side patch. The script then
checks that greedy 1-token outputs match between the dense and sparse vLLM
engines after the update.
The example performs the following steps:
* Load a training model on one GPU via a Ray actor.
* Launch a vLLM engine with the same real model on a second GPU.
* Verify trainer vs vLLM baseline agreement before any update.
* Apply a deterministic patch to ``model.embed_tokens.weight`` on the trainer.
* Run a dense NCCL update into a fresh vLLM engine and collect post-update
outputs.
* Reset the trainer back to the baseline checkpoint.
* Apply the same deterministic patch again.
* Run a sparse NCCL update into another fresh vLLM engine and collect
post-update outputs.
* Compare dense vs sparse baseline outputs, dense vs sparse post-update
outputs, estimated payload sizes, and trainer-side send times.
Current sparse weight transfer MVP limitations:
* ``TP=1`` and ``PP=1`` only
* sparse updates use runtime/kernel-format parameter names
* sparse updates are not composable with checkpoint-format or packed updates
This example assumes a single-node cluster with two GPUs.
"""
import hashlib
import os
import time
from collections.abc import Sequence
import ray
import torch
from ray.util.placement_group import placement_group
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
from transformers import AutoModelForCausalLM, AutoTokenizer
from vllm import LLM, SamplingParams
from vllm.config import WeightTransferConfig
from vllm.distributed.weight_transfer.base import SparseWeightPatch
from vllm.distributed.weight_transfer.nccl_engine import (
NCCLTrainerSendWeightsArgs,
NCCLWeightTransferEngine,
)
from vllm.utils.network_utils import get_ip, get_open_port
MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"
PATCHED_PARAM_NAME = "model.embed_tokens.weight"
MAX_PATCH_ROWS = 32
PROMPTS = [
"Hello, my name is",
"The president of the United States is",
"The capital of France is",
"The future of AI is",
]
SAMPLING_PARAMS = SamplingParams(temperature=0.0, max_tokens=1)
class MyLLM(LLM):
"""Configure the vLLM worker for Ray placement group execution."""
def __init__(self, *args, **kwargs):
os.environ["VLLM_RAY_BUNDLE_INDICES"] = "0"
super().__init__(*args, **kwargs)
@ray.remote(num_gpus=1)
class TrainModel:
"""Ray actor that owns the trainer-side model and deterministic patch state."""
def __init__(self, model_name: str):
self.model_name = model_name
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
if self.tokenizer.pad_token_id is None:
self.tokenizer.pad_token = self.tokenizer.eos_token
self.model = None
self.patched_param = None
self.pending_sparse_patches: list[SparseWeightPatch] | None = None
self.model_update_group = None
self.master_address = get_ip()
self.port = get_open_port()
self.reset_model()
def reset_model(self) -> None:
self.model = AutoModelForCausalLM.from_pretrained(
self.model_name,
torch_dtype=torch.bfloat16,
).to("cuda:0")
self.model.eval()
try:
self.patched_param = self.model.get_parameter(PATCHED_PARAM_NAME)
except AttributeError as exc:
raise RuntimeError(
f"Expected trainer model to expose `{PATCHED_PARAM_NAME}`"
) from exc
self.pending_sparse_patches = None
def create_rendezvous(self) -> tuple[str, int]:
self.port = get_open_port()
return self.master_address, self.port
def init_weight_transfer_group(self, world_size: int) -> None:
self.model_update_group = NCCLWeightTransferEngine.trainer_init(
dict(
master_address=self.master_address,
master_port=self.port,
world_size=world_size,
)
)
def get_dense_update_info(self, packed: bool = False) -> tuple[dict, int]:
names = []
dtype_names = []
shapes = []
payload_bytes = 0
for name, param in self.model.named_parameters():
names.append(name)
dtype_names.append(str(param.dtype).split(".")[-1])
shapes.append(list(param.shape))
payload_bytes += param.numel() * param.element_size()
return (
dict(
names=names,
dtype_names=dtype_names,
shapes=shapes,
packed=packed,
),
payload_bytes,
)
@torch.inference_mode()
def generate(
self,
prompts: Sequence[str],
max_new_tokens: int = 1,
) -> list[dict[str, object]]:
generations = []
for prompt in prompts:
model_inputs = self.tokenizer(prompt, return_tensors="pt").to("cuda:0")
output = self.model.generate(
**model_inputs,
max_new_tokens=max_new_tokens,
do_sample=False,
pad_token_id=self.tokenizer.pad_token_id,
)
new_token_ids = output[0, model_inputs["input_ids"].shape[1] :].tolist()
generations.append(
{
"token_ids": new_token_ids,
"text": self.tokenizer.decode(
new_token_ids,
skip_special_tokens=False,
),
}
)
return generations
def prepare_sparse_patch(
self,
prompts: Sequence[str],
max_patch_rows: int = MAX_PATCH_ROWS,
) -> tuple[dict[str, object], list[int], str, int]:
selected_token_ids: list[int] = []
special_ids = set(self.tokenizer.all_special_ids)
for prompt in prompts:
token_ids = self.tokenizer(prompt, add_special_tokens=False)["input_ids"]
for token_id in token_ids:
if token_id in special_ids or token_id in selected_token_ids:
continue
selected_token_ids.append(token_id)
if len(selected_token_ids) == max_patch_rows:
break
if len(selected_token_ids) == max_patch_rows:
break
if not selected_token_ids:
raise ValueError("Could not derive any non-special token IDs to patch")
vocab_size = self.patched_param.shape[0]
next_token_id = selected_token_ids[-1]
while len(selected_token_ids) < max_patch_rows:
next_token_id = (next_token_id + 1) % vocab_size
if next_token_id in special_ids or next_token_id in selected_token_ids:
continue
selected_token_ids.append(next_token_id)
row_ids = torch.tensor(
selected_token_ids,
device=self.patched_param.device,
dtype=torch.long,
)
hidden_size = self.patched_param.shape[1]
column_offsets = torch.arange(
hidden_size,
device=self.patched_param.device,
dtype=torch.long,
)
with torch.no_grad():
# Rotate the selected embedding rows instead of zeroing them so the
# patch remains deterministic while avoiding a degenerate collapse
# to the same special token after the update.
replacement_rows = self.patched_param[row_ids].roll(shifts=1, dims=0)
self.patched_param[row_ids] = replacement_rows
flat_indices = (
row_ids.unsqueeze(1).mul(hidden_size).add(column_offsets).reshape(-1)
)
flat_values = self.patched_param[row_ids].reshape(-1).contiguous()
self.pending_sparse_patches = [
SparseWeightPatch(
name=PATCHED_PARAM_NAME,
indices=flat_indices.to(torch.int32),
values=flat_values,
)
]
patch_digest = hashlib.sha256(
self.pending_sparse_patches[0].indices.cpu().numpy().tobytes()
+ self.pending_sparse_patches[0]
.values.detach()
.float()
.cpu()
.numpy()
.tobytes()
).hexdigest()
sparse_payload_bytes = (
flat_indices.numel() * torch.tensor([], dtype=torch.int32).element_size()
+ flat_values.numel() * flat_values.element_size()
)
update_info = dict(
names=[PATCHED_PARAM_NAME],
dtype_names=[str(self.patched_param.dtype).split(".")[-1]],
shapes=[list(self.patched_param.shape)],
num_updates_list=[flat_indices.numel()],
update_kind="sparse_flat",
)
return update_info, selected_token_ids, patch_digest, sparse_payload_bytes
def broadcast_weights(self, packed: bool = False) -> float:
if self.model_update_group is None:
raise RuntimeError("Weight transfer group is not initialized")
trainer_args = NCCLTrainerSendWeightsArgs(
group=self.model_update_group,
packed=packed,
)
start = time.perf_counter()
NCCLWeightTransferEngine.trainer_send_weights(
iterator=self.model.named_parameters(),
trainer_args=trainer_args,
)
torch.accelerator.synchronize()
return (time.perf_counter() - start) * 1000.0
def broadcast_pending_sparse_patch(self) -> float:
if self.model_update_group is None:
raise RuntimeError("Weight transfer group is not initialized")
if self.pending_sparse_patches is None:
raise RuntimeError("Sparse patch has not been prepared")
start = time.perf_counter()
NCCLWeightTransferEngine.trainer_send_sparse_weights(
iter(self.pending_sparse_patches),
NCCLTrainerSendWeightsArgs(group=self.model_update_group),
)
torch.accelerator.synchronize()
self.pending_sparse_patches = None
return (time.perf_counter() - start) * 1000.0
def launch_llm(
scheduling_inference: PlacementGroupSchedulingStrategy,
):
return ray.remote(
num_cpus=0,
num_gpus=0,
scheduling_strategy=scheduling_inference,
)(MyLLM).remote(
model=MODEL_NAME,
enforce_eager=True,
tensor_parallel_size=1,
distributed_executor_backend="ray",
gpu_memory_utilization=0.7,
weight_transfer_config=WeightTransferConfig(backend="nccl"),
)
def collect_vllm_generations(llm_handle) -> list[dict[str, object]]:
outputs = ray.get(llm_handle.generate.remote(PROMPTS, SAMPLING_PARAMS))
generations = []
for output in outputs:
generations.append(
{
"token_ids": output.outputs[0].token_ids,
"text": output.outputs[0].text,
}
)
return generations
def token_sequences_match(
left: Sequence[dict[str, object]],
right: Sequence[dict[str, object]],
) -> bool:
return [item["token_ids"] for item in left] == [item["token_ids"] for item in right]
def print_generations(label: str, prompts: Sequence[str], generations) -> None:
print(f"\n{label}")
print("-" * 50)
for prompt, generation in zip(prompts, generations):
print(f"Prompt: {prompt!r}")
print(f"Token IDs: {generation['token_ids']}")
print(f"Text: {generation['text']!r}")
print("-" * 50)
def run_dense_phase(
train_model,
scheduling_inference: PlacementGroupSchedulingStrategy,
) -> dict[str, object]:
ray.get(train_model.reset_model.remote())
llm = launch_llm(scheduling_inference)
try:
dense_before = collect_vllm_generations(llm)
ray.get(llm.sleep.remote(level=0))
master_address, master_port = ray.get(train_model.create_rendezvous.remote())
world_size = ray.get(llm.get_world_size.remote()) + 1
inference_init = llm.init_weight_transfer_engine.remote(
dict(
init_info=dict(
master_address=master_address,
master_port=master_port,
rank_offset=1,
world_size=world_size,
)
)
)
trainer_init = train_model.init_weight_transfer_group.remote(world_size)
ray.get([trainer_init, inference_init])
ray.get(llm.start_weight_update.remote(is_checkpoint_format=True))
dense_update_info, dense_payload_bytes = ray.get(
train_model.get_dense_update_info.remote()
)
_, selected_token_ids, patch_digest, _ = ray.get(
train_model.prepare_sparse_patch.remote(PROMPTS)
)
inference_update = llm.update_weights.remote(
dict(update_info=dense_update_info)
)
dense_send_ms, _ = ray.get(
[
train_model.broadcast_weights.remote(packed=False),
inference_update,
]
)
ray.get(llm.finish_weight_update.remote())
ray.get(llm.wake_up.remote(tags=["scheduling"]))
dense_after = collect_vllm_generations(llm)
return {
"dense_before": dense_before,
"dense_after": dense_after,
"selected_token_ids": selected_token_ids,
"patch_digest": patch_digest,
"dense_payload_bytes": dense_payload_bytes,
"dense_send_ms": dense_send_ms,
}
finally:
ray.kill(llm)
def run_sparse_phase(
train_model,
scheduling_inference: PlacementGroupSchedulingStrategy,
) -> dict[str, object]:
ray.get(train_model.reset_model.remote())
llm = launch_llm(scheduling_inference)
try:
sparse_before = collect_vllm_generations(llm)
ray.get(llm.sleep.remote(level=0))
master_address, master_port = ray.get(train_model.create_rendezvous.remote())
world_size = ray.get(llm.get_world_size.remote()) + 1
inference_init = llm.init_weight_transfer_engine.remote(
dict(
init_info=dict(
master_address=master_address,
master_port=master_port,
rank_offset=1,
world_size=world_size,
)
)
)
trainer_init = train_model.init_weight_transfer_group.remote(world_size)
ray.get([trainer_init, inference_init])
ray.get(llm.start_weight_update.remote(is_checkpoint_format=False))
sparse_update_info, selected_token_ids, patch_digest, sparse_payload_bytes = (
ray.get(train_model.prepare_sparse_patch.remote(PROMPTS))
)
inference_update = llm.update_weights.remote(
dict(update_info=sparse_update_info)
)
sparse_send_ms, _ = ray.get(
[
train_model.broadcast_pending_sparse_patch.remote(),
inference_update,
]
)
ray.get(llm.finish_weight_update.remote())
ray.get(llm.wake_up.remote(tags=["scheduling"]))
sparse_after = collect_vllm_generations(llm)
return {
"sparse_before": sparse_before,
"sparse_after": sparse_after,
"selected_token_ids": selected_token_ids,
"patch_digest": patch_digest,
"sparse_payload_bytes": sparse_payload_bytes,
"sparse_send_ms": sparse_send_ms,
}
finally:
ray.kill(llm)
ray.init()
try:
train_model = TrainModel.remote(MODEL_NAME)
pg_inference = placement_group([{"GPU": 1, "CPU": 0}])
ray.get(pg_inference.ready())
scheduling_inference = PlacementGroupSchedulingStrategy(
placement_group=pg_inference,
placement_group_capture_child_tasks=True,
placement_group_bundle_index=0,
)
dense_results = run_dense_phase(train_model, scheduling_inference)
sparse_results = run_sparse_phase(train_model, scheduling_inference)
baseline_equal = token_sequences_match(
dense_results["dense_before"],
sparse_results["sparse_before"],
)
patch_selection_equal = (
dense_results["selected_token_ids"] == sparse_results["selected_token_ids"]
)
patch_digest_equal = dense_results["patch_digest"] == sparse_results["patch_digest"]
after_equal = token_sequences_match(
dense_results["dense_after"],
sparse_results["sparse_after"],
)
any_output_changed = any(
before["token_ids"] != after["token_ids"]
for before, after in zip(
dense_results["dense_before"],
dense_results["dense_after"],
)
)
dense_payload_mb = dense_results["dense_payload_bytes"] / (1024 * 1024)
sparse_payload_mb = sparse_results["sparse_payload_bytes"] / (1024 * 1024)
print_generations(
"Dense baseline outputs",
PROMPTS,
dense_results["dense_before"],
)
print_generations(
"Sparse baseline outputs", PROMPTS, sparse_results["sparse_before"]
)
print_generations(
"Dense outputs after update", PROMPTS, dense_results["dense_after"]
)
print_generations(
"Sparse outputs after update",
PROMPTS,
sparse_results["sparse_after"],
)
print(f"patched_token_ids = {dense_results['selected_token_ids']}")
print(f"patch_selection_equal = {patch_selection_equal}")
print(f"dense_patch_digest = {dense_results['patch_digest']}")
print(f"sparse_patch_digest = {sparse_results['patch_digest']}")
print(f"patch_digest_equal = {patch_digest_equal}")
print(f"baseline_equal = {baseline_equal}")
print(f"after_equal = {after_equal}")
print(f"any_output_changed = {any_output_changed}")
print(f"dense_payload_mb = {dense_payload_mb:.2f}")
print(f"sparse_payload_mb = {sparse_payload_mb:.2f}")
print(f"dense_send_ms = {dense_results['dense_send_ms']:.2f}")
print(f"sparse_send_ms = {sparse_results['sparse_send_ms']:.2f}")
if not baseline_equal:
raise RuntimeError(
"Dense and sparse phases did not start from the same baseline"
)
if not patch_selection_equal:
raise RuntimeError("Dense and sparse phases used different sparse patches")
if not patch_digest_equal:
raise RuntimeError("Dense and sparse phases produced different patch values")
if not after_equal:
raise RuntimeError("Dense and sparse updates produced different outputs")
if not any_output_changed:
raise RuntimeError("Patch did not change the observed outputs")
finally:
ray.shutdown()
+5
View File
@@ -162,7 +162,12 @@ dout = "dout"
Pn = "Pn"
arange = "arange"
thw = "thw"
# temporal position ids (parallels hpos/wpos in vision RoPE)
tpos = "tpos"
subtile = "subtile"
subtiles = "subtiles"
reord = "reord"
Ot = "Ot"
HSA = "HSA"
setp = "setp"
CPY = "CPY"
+1
View File
@@ -29,6 +29,7 @@ xgrammar >= 0.2.0, < 1.0.0; platform_machine == "x86_64" or platform_machine ==
typing_extensions >= 4.10
filelock >= 3.16.1 # need to contain https://github.com/tox-dev/filelock/pull/317
partial-json-parser # used for parsing partial JSON outputs
jsonschema >= 4.23.0 # required for MiniMax M3 tool schema validation
pyzmq >= 25.0.0
msgspec
gguf >= 0.17.0
+2 -2
View File
@@ -9,8 +9,8 @@ torchaudio==2.11.0
# These must be updated alongside torch
torchvision==0.26.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
# FlashInfer should be updated together with the Dockerfile
flashinfer-python==0.6.11.post2
flashinfer-cubin==0.6.11.post2
flashinfer-python==0.6.12
flashinfer-cubin==0.6.12
apache-tvm-ffi==0.1.9
tilelang==0.1.9
# Cap nvidia-cudnn-frontend (transitive dep of flashinfer) due to
+1
View File
@@ -360,6 +360,7 @@ jsonpointer==3.0.0
# via jsonschema
jsonschema==4.23.0
# via
# -c requirements/common.txt
# hypothesis-jsonschema
# mistral-common
# ray
+2
View File
@@ -440,6 +440,8 @@ jsonpointer==3.1.0
# via jsonschema
jsonschema==4.26.0
# via
# -c requirements/common.txt
# -r requirements/test/../common.txt
# hypothesis-jsonschema
# mcp
# mistral-common
+1
View File
@@ -229,6 +229,7 @@ jsonlines==4.0.0
# via lm-eval
jsonschema==4.26.0
# via
# -c requirements/common.txt
# hypothesis-jsonschema
# mistral-common
# schemathesis
+87
View File
@@ -3458,6 +3458,75 @@ version = "0.1.29"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e0c5ccf5294c6ccd63a74f1565028353830a9c2f5eb0c682c355c471726a6e3f"
[[package]]
name = "pyo3"
version = "0.28.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "91fd8e38a3b50ed1167fb981cd6fd60147e091784c427b8f7183a7ee32c31c12"
dependencies = [
"libc",
"once_cell",
"portable-atomic",
"pyo3-build-config",
"pyo3-ffi",
"pyo3-macros",
]
[[package]]
name = "pyo3-build-config"
version = "0.28.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e368e7ddfdeb98c9bca7f8383be1648fd84ab466bf2bc015e94008db6d35611e"
dependencies = [
"target-lexicon",
]
[[package]]
name = "pyo3-ffi"
version = "0.28.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7f29e10af80b1f7ccaf7f69eace800a03ecd13e883acfacc1e5d0988605f651e"
dependencies = [
"libc",
"pyo3-build-config",
]
[[package]]
name = "pyo3-macros"
version = "0.28.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "df6e520eff47c45997d2fc7dd8214b25dd1310918bbb2642156ef66a67f29813"
dependencies = [
"proc-macro2",
"pyo3-macros-backend",
"quote",
"syn 2.0.117",
]
[[package]]
name = "pyo3-macros-backend"
version = "0.28.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c4cdc218d835738f81c2338f822078af45b4afdf8b2e33cbb5916f108b813acb"
dependencies = [
"heck",
"proc-macro2",
"pyo3-build-config",
"quote",
"syn 2.0.117",
]
[[package]]
name = "pythonize"
version = "0.28.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0b79f670c9626c8b651c0581011b57b6ba6970bb69faf01a7c4c0cfc81c43f95"
dependencies = [
"pyo3",
"serde",
"serde_json",
]
[[package]]
name = "qoi"
version = "0.4.1"
@@ -4669,6 +4738,12 @@ dependencies = [
"libc",
]
[[package]]
name = "target-lexicon"
version = "0.13.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca"
[[package]]
name = "task-local"
version = "0.1.1"
@@ -5622,6 +5697,7 @@ dependencies = [
"expect-test",
"futures",
"half",
"indexmap 2.13.0",
"itertools 0.14.0",
"llm-multimodal",
"minijinja",
@@ -5900,6 +5976,17 @@ dependencies = [
"winnow",
]
[[package]]
name = "vllm-tool-parser-py"
version = "0.1.0"
dependencies = [
"pyo3",
"pythonize",
"serde_json",
"thiserror-ext",
"vllm-tool-parser",
]
[[package]]
name = "walkdir"
version = "2.5.0"
+5 -1
View File
@@ -12,6 +12,7 @@ members = [
"src/text",
"src/tokenizer",
"src/tool-parser",
"src/tool-parser/python",
]
resolver = "3"
@@ -43,6 +44,7 @@ half = { version = "2.7.1", features = ["bytemuck"] }
hex = "0.4.3"
hf-hub = { version = "0.5.0", features = ["tokio"] }
http-body = "1.0.1"
indexmap = "2.13.0"
itertools = "0.14.0"
libc = "0.2.177"
llm-multimodal = { git = "https://github.com/vllm-project/llm-multimodal", rev = "5b558989844d1c7af3e43d0f604069ffd9c06320" }
@@ -59,6 +61,8 @@ prometheus-client = "0.24.0"
prometheus-client-derive-encode = "0.5.0"
prost = "0.14.3"
prost-types = "0.14.3"
pyo3 = "0.28.3"
pythonize = "0.28.0"
rand = "0.9.2"
reasoning-parser = "1.2.2"
reqwest = { version = "0.12.8", default-features = false, features = ["rustls-tls"] }
@@ -69,7 +73,7 @@ rustc-hash = "1.1.0"
serde = { version = "1.0.228", features = ["derive"] }
serde-json-fmt = "0.1.0"
serde_default = "0.2.0"
serde_json = { version = "1.0.145", features = ["arbitrary_precision", "preserve_order"] }
serde_json = { version = "1.0.145", features = ["preserve_order"] }
serde_repr = "0.1.20"
serde_tuple = "1.1.3"
serde_with = "3.18.0"
+1
View File
@@ -10,6 +10,7 @@ asynk-strim-attr.workspace = true
easy-ext.workspace = true
futures.workspace = true
half.workspace = true
indexmap.workspace = true
itertools.workspace = true
llm-multimodal.workspace = true
minijinja.workspace = true
+1 -1
View File
@@ -233,7 +233,7 @@ mod tests {
)
.unwrap_err();
expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, hermes, hy_v3, internlm, kimi_k2, llama3_json, llama4_json, minimax_m2, mistral, qwen3_coder, qwen3_xml)"].assert_eq(&error.to_report_string());
expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, hermes, hy_v3, kimi_k2, llama3_json, llama4_json, minimax_m2, minimax_m3, mistral, qwen3_coder, qwen3_xml)"].assert_eq(&error.to_report_string());
}
#[test]
+7 -2
View File
@@ -5,8 +5,9 @@ use std::sync::LazyLock;
pub use vllm_reasoning_parser::{
CohereCmdReasoningParser, DeepSeekR1ReasoningParser, DeepSeekV3ReasoningParser,
DeepSeekV4ReasoningParser, Gemma4ReasoningParser, Glm45ReasoningParser, KimiK2ReasoningParser,
KimiReasoningParser, MiniMaxM2ReasoningParser, NemotronV3ReasoningParser, Qwen3ReasoningParser,
ReasoningDelta, ReasoningError, ReasoningParser, Step3ReasoningParser,
KimiReasoningParser, MiniMaxM2ReasoningParser, MiniMaxM3ReasoningParser,
NemotronV3ReasoningParser, Qwen3ReasoningParser, ReasoningDelta, ReasoningError,
ReasoningParser, Step3ReasoningParser,
};
use vllm_tokenizer::DynTokenizer;
@@ -23,6 +24,7 @@ pub mod names {
pub const KIMI: &str = "kimi";
pub const KIMI_K2: &str = "kimi_k2";
pub const MINIMAX_M2: &str = "minimax_m2";
pub const MINIMAX_M3: &str = "minimax_m3";
pub const NEMOTRON_V3: &str = "nemotron_v3";
pub const QWEN3: &str = "qwen3";
pub const STEP3: &str = "step3";
@@ -59,6 +61,7 @@ impl ReasoningParserFactory {
.register_parser::<KimiReasoningParser>(names::KIMI)
.register_parser::<KimiK2ReasoningParser>(names::KIMI_K2)
.register_parser::<MiniMaxM2ReasoningParser>(names::MINIMAX_M2)
.register_parser::<MiniMaxM3ReasoningParser>(names::MINIMAX_M3)
.register_parser::<NemotronV3ReasoningParser>(names::NEMOTRON_V3)
.register_parser::<Qwen3ReasoningParser>(names::QWEN3)
.register_parser::<Step3ReasoningParser>(names::STEP3);
@@ -78,6 +81,8 @@ impl ReasoningParserFactory {
.register_pattern("kimi-k2", names::KIMI_K2)
.register_pattern("kimi", names::KIMI)
.register_pattern("step3", names::STEP3)
.register_pattern("minimax-m3", names::MINIMAX_M3)
.register_pattern("mm-m3", names::MINIMAX_M3)
.register_pattern("minimax", names::MINIMAX_M2)
.register_pattern("mm-m2", names::MINIMAX_M2)
.register_pattern("cohere", names::COHERE_CMD)
@@ -32,8 +32,10 @@ fn factory_contains_and_lists_registered_parsers() {
let factory = ReasoningParserFactory::new();
assert!(factory.contains(names::QWEN3));
assert!(factory.contains(names::DEEPSEEK_V4));
assert!(factory.contains(names::MINIMAX_M3));
assert!(factory.list().contains(&names::QWEN3.to_string()));
assert!(factory.list().contains(&names::DEEPSEEK_V4.to_string()));
assert!(factory.list().contains(&names::MINIMAX_M3.to_string()));
}
#[test]
@@ -49,6 +51,19 @@ fn factory_resolves_deepseek_v4_to_qwen3_alias() {
);
}
#[test]
fn factory_resolves_minimax_m3_before_generic_minimax() {
let factory = ReasoningParserFactory::new();
assert_eq!(
factory.resolve_name_for_model("MiniMaxAI/Minimax-M3-preview"),
Some(names::MINIMAX_M3)
);
assert_eq!(
factory.resolve_name_for_model("mm-m3"),
Some(names::MINIMAX_M3)
);
}
#[test]
fn factory_rejects_unknown_parser_names() {
let tokenizer = Arc::new(FakeTokenizer);
+5 -11
View File
@@ -5,7 +5,7 @@ use std::sync::LazyLock;
pub use vllm_tool_parser::{
DeepSeekV3ToolParser, DeepSeekV4ToolParser, DeepSeekV31ToolParser, DeepSeekV32ToolParser,
Gemma4ToolParser, Glm45MoeToolParser, Glm47MoeToolParser, HermesToolParser, HyV3ToolParser,
Internlm2ToolParser, KimiK2ToolParser, Llama3JsonToolParser, MinimaxM2ToolParser,
KimiK2ToolParser, Llama3JsonToolParser, MinimaxM2ToolParser, MinimaxM3ToolParser,
MistralToolParser, Qwen3CoderToolParser, Qwen3XmlToolParser, ToolCallDelta, ToolParser,
ToolParserError, ToolParserOutput,
};
@@ -24,13 +24,11 @@ pub mod names {
pub const GEMMA4: &str = "gemma4";
pub const HERMES: &str = "hermes";
pub const HY_V3: &str = "hy_v3";
// Matches the Python CLI name `--tool-call-parser internlm`, which Python
// also routes to `Internlm2ToolParser` despite the version-agnostic name.
pub const INTERNLM: &str = "internlm";
pub const KIMI_K2: &str = "kimi_k2";
pub const LLAMA3_JSON: &str = "llama3_json";
pub const LLAMA4_JSON: &str = "llama4_json";
pub const MINIMAX_M2: &str = "minimax_m2";
pub const MINIMAX_M3: &str = "minimax_m3";
pub const MISTRAL: &str = "mistral";
pub const QWEN3_CODER: &str = "qwen3_coder";
pub const QWEN3_XML: &str = "qwen3_xml";
@@ -65,11 +63,11 @@ impl ToolParserFactory {
.register_parser::<Gemma4ToolParser>(names::GEMMA4)
.register_parser::<HermesToolParser>(names::HERMES)
.register_parser::<HyV3ToolParser>(names::HY_V3)
.register_parser::<Internlm2ToolParser>(names::INTERNLM)
.register_parser::<KimiK2ToolParser>(names::KIMI_K2)
.register_parser::<Llama3JsonToolParser>(names::LLAMA3_JSON)
.register_parser::<Llama3JsonToolParser>(names::LLAMA4_JSON)
.register_parser::<MinimaxM2ToolParser>(names::MINIMAX_M2)
.register_parser::<MinimaxM3ToolParser>(names::MINIMAX_M3)
.register_parser::<MistralToolParser>(names::MISTRAL)
.register_parser::<Qwen3XmlToolParser>(names::QWEN3_XML)
.register_parser::<Qwen3CoderToolParser>(names::QWEN3_CODER);
@@ -84,12 +82,6 @@ impl ToolParserFactory {
.register_pattern("hermes", names::HERMES)
.register_pattern("hy3", names::HY_V3)
.register_pattern("hy_v3", names::HY_V3)
// Narrow to `internlm2` substring so it matches `internlm2-chat-7b`
// and `internlm2_5-7b-chat` but NOT `internlm-chat-7b` (InternLM v1,
// routes to Llama), `internlm3-*` (also Llama-architecture per
// vllm/model_executor/models/registry.py:146), or `Intern-S1` /
// `Intern-S1-Pro` (separate intern-s1 parser, see PR #40115).
.register_pattern("internlm2", names::INTERNLM)
.register_pattern("llama-4", names::LLAMA4_JSON)
.register_pattern("llama-3.2", names::LLAMA3_JSON)
.register_pattern("llama-3.1", names::LLAMA3_JSON)
@@ -106,6 +98,8 @@ impl ToolParserFactory {
.register_pattern("gemma4", names::GEMMA4)
.register_pattern("gemma-4", names::GEMMA4)
.register_pattern("kimi-k2", names::KIMI_K2)
.register_pattern("minimax-m3", names::MINIMAX_M3)
.register_pattern("mm-m3", names::MINIMAX_M3)
.register_pattern("minimax", names::MINIMAX_M2)
.register_pattern("mm-m2", names::MINIMAX_M2);
+8 -29
View File
@@ -153,6 +153,14 @@ fn factory_new_resolves_default_patterns() {
factory.resolve_name_for_model("tencent/Hy3-preview"),
Some(names::HY_V3)
);
assert_eq!(
factory.resolve_name_for_model("MiniMax/MiniMax-M3-Text"),
Some(names::MINIMAX_M3)
);
assert_eq!(
factory.resolve_name_for_model("org/mm-m3-base"),
Some(names::MINIMAX_M3)
);
assert_eq!(
factory.resolve_name_for_model("MiniMax/MiniMax-M2-01"),
Some(names::MINIMAX_M2)
@@ -161,33 +169,4 @@ fn factory_new_resolves_default_patterns() {
factory.resolve_name_for_model("org/mm-m2-base"),
Some(names::MINIMAX_M2)
);
// InternLM2 positive: both dashed and underscored versioned names route.
assert_eq!(
factory.resolve_name_for_model("internlm/internlm2-chat-7b"),
Some(names::INTERNLM)
);
assert_eq!(
factory.resolve_name_for_model("internlm/internlm2_5-7b-chat"),
Some(names::INTERNLM)
);
// Negative: other internlm-org models do NOT route to the InternLM2 parser,
// since they use unrelated prompt formats.
// - InternLM v1 (`internlm-chat-7b`) routes to Llama
// - InternLM3 (`internlm3-8b-instruct`) routes to Llama
// - Intern-S1 / Intern-S1-Pro have their own parser (Python PR #40115)
assert_eq!(
factory.resolve_name_for_model("internlm/internlm-chat-7b"),
None
);
assert_eq!(
factory.resolve_name_for_model("internlm/internlm3-8b-instruct"),
None
);
assert_eq!(factory.resolve_name_for_model("internlm/Intern-S1"), None);
assert_eq!(
factory.resolve_name_for_model("internlm/Intern-S1-Pro"),
None
);
}
+33 -7
View File
@@ -1,7 +1,7 @@
use std::collections::HashMap;
use serde::Serialize;
use serde_json::Value;
use serde_json::Value as JsonValue;
use thiserror_ext::AsReport as _;
use tracing::{info, trace, warn};
use vllm_text::Prompt;
@@ -13,6 +13,7 @@ use self::format::{
ChatTemplateContentFormat, ChatTemplateContentFormatOption as ContentFormatOption,
};
use self::template::{CompiledChatTemplate, TemplateContext};
use self::value::{TemplateValue, to_template_value};
use super::{ChatRenderer, RenderedPrompt};
use crate::error::Result;
use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest};
@@ -24,6 +25,7 @@ mod error;
mod format;
mod template;
mod tojson;
mod value;
pub use template::{load_chat_template, resolve_chat_template};
@@ -38,7 +40,7 @@ pub struct MultimodalRenderInfo {
/// state.
pub struct HfChatRenderer {
default_template: Option<CompiledChatTemplate>,
default_template_kwargs: HashMap<String, Value>,
default_template_kwargs: HashMap<String, JsonValue>,
content_format: ContentFormatOption,
special_tokens: Option<HfSpecialTokens>,
multimodal: Option<MultimodalRenderInfo>,
@@ -48,7 +50,7 @@ impl HfChatRenderer {
/// Create a renderer from the given template string.
pub fn new(
template: Option<String>,
default_template_kwargs: HashMap<String, Value>,
default_template_kwargs: HashMap<String, JsonValue>,
content_format: ContentFormatOption,
) -> Result<Self> {
Ok(Self {
@@ -245,7 +247,7 @@ struct TemplateToolCall {
#[derive(Debug, Serialize)]
struct TemplateToolFunction {
name: String,
arguments: Value,
arguments: TemplateValue,
}
#[derive(Debug, Serialize)]
@@ -259,7 +261,7 @@ pub(super) struct TemplateTool {
struct TemplateToolDefinition {
name: String,
description: Option<String>,
parameters: Value,
parameters: TemplateValue,
strict: Option<bool>,
}
@@ -345,13 +347,14 @@ fn to_template_tool_calls(
let mut tool_calls = Vec::new();
for tool_call in content.tool_calls() {
let arguments = serde_json::from_str::<Value>(&tool_call.arguments).map_err(|error| {
let arguments = serde_json::from_str(&tool_call.arguments).map_err(|error| {
Error::ChatTemplate(format!(
"assistant tool call `{}` has invalid JSON arguments: {}",
tool_call.id,
error.as_report()
))
})?;
let arguments = to_template_value(arguments);
tool_calls.push(TemplateToolCall {
id: tool_call.id.clone(),
@@ -434,7 +437,7 @@ fn to_template_tools(tools: &[ChatTool]) -> Vec<TemplateTool> {
function: TemplateToolDefinition {
name: tool.name.clone(),
description: tool.description.clone(),
parameters: tool.parameters.clone(),
parameters: to_template_value(tool.parameters.clone()),
strict: tool.strict,
},
})
@@ -909,6 +912,29 @@ mod tests {
assert_eq!(rendered, "get_weather|Paris|call_1|Sunny");
}
#[test]
fn chat_template_tool_call_argument_items_method_is_not_shadowed_by_field() {
let request = sample_request(vec![ChatMessage::assistant_blocks(vec![
AssistantContentBlock::ToolCall(crate::AssistantToolCall {
id: "call_1".to_string(),
name: "add".to_string(),
arguments: r#"{"items":"operands","x":2,"y":1.0}"#.to_string(),
}),
])]);
let rendered = render(
Some(
"{%- set arguments = messages[0].tool_calls[0].function.arguments -%}
{%- for key, value in arguments.items() -%}{{ key }}={{ value }};{%- endfor -%}
|{{ arguments['items'] }}",
),
&request,
)
.unwrap();
assert_eq!(rendered, "items=operands;x=2;y=1.0;|operands");
}
#[test]
fn qwen35_template_renders_prefilled_reasoning_start_when_thinking_enabled() {
let mut request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]);
+18 -2
View File
@@ -208,11 +208,27 @@ mod tests {
}
#[test]
fn tojson_preserves_arbitrary_precision_number_spelling() {
fn tojson_uses_standard_serde_json_number_spelling() {
let payload = serde_json::from_str(r#"{"x":2,"y":1.00}"#).unwrap();
let rendered = render("{{ payload|tojson }}", payload);
assert_eq!(rendered, "{\"x\": 2, \"y\": 1.00}");
// TODO: we cannot preserve the original number precision by enabling `serde_json`'s
// `arbitrary_precision` feature, otherwise the following test
// `serialized_json_numbers_do_not_leak_serde_private_representation` will fail.
// See issue: https://github.com/mitsuhiko/minijinja/issues/641
assert_eq!(rendered, "{\"x\": 2, \"y\": 1.0}");
}
#[test]
fn serialized_json_numbers_do_not_leak_serde_private_representation() {
let payload: serde_json::Value = serde_json::from_str(r#"{"x":2,"y":1.00}"#).unwrap();
let rendered = render("{{ payload }}", payload);
// TODO: we cannot preserve the original number precision by enabling `serde_json`'s
// `arbitrary_precision` feature, otherwise this will fail.
// See issue: https://github.com/mitsuhiko/minijinja/issues/641
assert!(!rendered.contains("$serde_json::private::Number"));
assert_eq!(rendered, r#"{"x": 2, "y": 1.0}"#);
}
#[test]
+77
View File
@@ -0,0 +1,77 @@
use std::sync::Arc;
use indexmap::IndexMap;
use minijinja::value::{Enumerator, Object, ObjectExt, ObjectRepr};
use minijinja::{Error as TemplateError, ErrorKind as TemplateErrorKind, State};
use serde::Serialize;
use serde_json::Value as JsonValue;
/// A wrapper around `minijinja::Value` that can be constructed with `to_template_value` and used
/// as a value in the chat template.
#[derive(Debug, Serialize)]
#[serde(transparent)]
pub(super) struct TemplateValue(minijinja::Value);
pub(super) fn to_template_value(value: JsonValue) -> TemplateValue {
TemplateValue(match value {
JsonValue::Array(values) => values
.into_iter()
.map(to_template_value)
.map(|value| value.0)
.collect::<minijinja::Value>(),
JsonValue::Object(values) => minijinja::Value::from_object(TemplateMap(
values
.into_iter()
.map(|(key, value)| (key, to_template_value(value).0))
.collect(),
)),
// For primitive values, directly convert them to `minijinja::Value` using `from_serialize`.
value => minijinja::Value::from_serialize(value),
})
}
/// A custom map type that always returns `UnknownMethod` for method calls, so that pycompat can
/// always handle dict methods through the unknown-method callback.
///
/// Use `IndexMap` to preserve the original key order when iterating.
///
/// MiniJinja's default map can resolve a same-named field before Python dict methods. HF templates
/// commonly call `dict.items()`, which would fail if the map had an `items` field.
/// See issue: https://github.com/mitsuhiko/minijinja/issues/903
#[derive(Debug)]
struct TemplateMap(IndexMap<String, minijinja::Value>);
impl Object for TemplateMap {
fn repr(self: &Arc<Self>) -> ObjectRepr {
ObjectRepr::Map
}
fn get_value(self: &Arc<Self>, key: &minijinja::Value) -> Option<minijinja::Value> {
self.0.get(key.as_str()?).cloned()
}
fn get_value_by_str(self: &Arc<Self>, key: &str) -> Option<minijinja::Value> {
self.0.get(key).cloned()
}
fn enumerate(self: &Arc<Self>) -> Enumerator {
self.mapped_rev_enumerator(|this| {
Box::new(this.0.keys().map(|key| minijinja::Value::from(key.as_str())))
})
}
fn enumerator_len(self: &Arc<Self>) -> Option<usize> {
Some(self.0.len())
}
fn call_method(
self: &Arc<Self>,
_state: &State<'_, '_>,
_method: &str,
_args: &[minijinja::Value],
) -> std::result::Result<minijinja::Value, TemplateError> {
// Always return `UnknownMethod` for method calls,
// so that pycompat can handle dict methods through the unknown-method callback.
Err(TemplateError::from(TemplateErrorKind::UnknownMethod))
}
}
+33 -82
View File
@@ -20,7 +20,6 @@ use vllm_chat::{
use vllm_text::{DecodedTextEvent, Finished, Prompt};
/// One model/parser configuration used to run the fixed roundtrip fixtures.
#[derive(Clone)]
struct RoundtripCase {
/// Hugging Face model id resolved through the production backend loader.
model_id: &'static str,
@@ -32,45 +31,11 @@ struct RoundtripCase {
tool_call_parser: ParserSelection,
/// Reasoning parser selection used by the output processor.
reasoning_parser: ParserSelection,
/// How this model's chat template handles thinking mode.
thinking_behavior: ThinkingBehavior,
/// JSON formatting expected after this model's template has materialized
/// tool-call arguments.
json_fmt: JsonFmt,
}
#[derive(Clone, Copy)]
enum ThinkingBehavior {
/// The chat template accepts explicit thinking on/off kwargs, and uses
/// `default` when the request does not specify either kwarg.
Toggleable { default: bool },
/// The chat template always behaves as `value` for this fixture.
Always { value: bool },
}
impl ThinkingBehavior {
fn default(self) -> bool {
match self {
Self::Toggleable { default } => default,
Self::Always { value } => value,
}
}
fn fixtures(self) -> Vec<Option<bool>> {
match self {
Self::Toggleable { .. } => vec![
Some(true), // explicitly enable thinking
Some(false), // explicitly disable thinking
None, // use default template behavior
],
Self::Always { value } => vec![
Some(value), // explicitly request the supported thinking behavior
None, // use default template behavior
],
}
}
}
impl RoundtripCase {
/// Qwen3 XML tool-call format with `qwen3` reasoning tags.
fn qwen3() -> Self {
@@ -79,7 +44,6 @@ impl RoundtripCase {
assistant_stop_suffix: "<|im_end|>\n",
tool_call_parser: ParserSelection::Auto,
reasoning_parser: ParserSelection::Auto,
thinking_behavior: ThinkingBehavior::Toggleable { default: true },
json_fmt: spaced_json_fmt(),
}
}
@@ -91,7 +55,6 @@ impl RoundtripCase {
assistant_stop_suffix: "<|im_end|>\n",
tool_call_parser: ParserSelection::Auto,
reasoning_parser: ParserSelection::Auto,
thinking_behavior: ThinkingBehavior::Toggleable { default: true },
json_fmt: compact_json_fmt(),
}
}
@@ -103,7 +66,6 @@ impl RoundtripCase {
assistant_stop_suffix: "[e~[\n",
tool_call_parser: ParserSelection::Auto,
reasoning_parser: ParserSelection::Auto,
thinking_behavior: ThinkingBehavior::Always { value: true },
json_fmt: compact_json_fmt(),
}
}
@@ -115,7 +77,6 @@ impl RoundtripCase {
assistant_stop_suffix: "<end▁of▁sentence>",
tool_call_parser: ParserSelection::Auto,
reasoning_parser: ParserSelection::Auto,
thinking_behavior: ThinkingBehavior::Toggleable { default: false },
json_fmt: compact_json_fmt(),
}
}
@@ -127,7 +88,6 @@ impl RoundtripCase {
assistant_stop_suffix: "",
tool_call_parser: ParserSelection::Auto,
reasoning_parser: ParserSelection::Auto,
thinking_behavior: ThinkingBehavior::Toggleable { default: true },
json_fmt: compact_json_fmt(),
}
}
@@ -140,7 +100,6 @@ impl RoundtripCase {
assistant_stop_suffix: "<|im_end|>",
tool_call_parser: ParserSelection::Auto,
reasoning_parser: ParserSelection::Auto,
thinking_behavior: ThinkingBehavior::Toggleable { default: true },
json_fmt: spaced_json_fmt(),
}
}
@@ -176,44 +135,35 @@ roundtrip_tests! {
/// Run the fixed reasoning+content fixture for one model/parser case.
async fn run_roundtrip_reasoning_and_content(case: RoundtripCase) -> Result<()> {
for thinking in case.thinking_behavior.fixtures() {
run_roundtrip_reasoning_and_content_inner(case.clone(), thinking).await?;
}
Ok(())
}
async fn run_roundtrip_reasoning_and_content_inner(
case: RoundtripCase,
thinking: Option<bool>,
) -> Result<()> {
let backends = load_roundtrip_backends(&case).await?;
let request = roundtrip_request(
"roundtrip-reasoning-content",
vec![ChatMessage::text(ChatRole::User, "What is 2 + 2?")],
Vec::new(),
thinking,
);
let expected_reasoning = "Need compute 2 + 2 directly.";
let expected_text = "The answer is 4.";
let effective_thinking = thinking.unwrap_or(case.thinking_behavior.default());
let assistant = {
let mut content = Vec::new();
if effective_thinking {
content.push(AssistantContentBlock::Reasoning {
text: expected_reasoning.to_string(),
});
}
content.push(AssistantContentBlock::Text {
text: expected_text.to_string(),
});
AssistantMessage { content }
};
let result = run_roundtrip(&case, &backends, &request, assistant).await?;
let result = run_roundtrip(
&case,
&backends,
&request,
AssistantMessage {
content: vec![
AssistantContentBlock::Reasoning {
text: expected_reasoning.to_string(),
},
AssistantContentBlock::Text {
text: expected_text.to_string(),
},
],
},
)
.await?;
assert_eq!(
result.parsed_message.reasoning().as_deref().map(str::trim),
effective_thinking.then_some(expected_reasoning)
Some(expected_reasoning)
);
assert_eq!(result.parsed_message.text().trim(), expected_text);
assert_eq!(result.parsed_message.tool_calls().count(), 0);
@@ -233,10 +183,9 @@ async fn run_roundtrip_tool_call_mix(case: RoundtripCase) -> Result<()> {
"roundtrip-reasoning-tools",
vec![ChatMessage::text(
ChatRole::User,
"Check Shanghai weather and add 1.00 plus 2.",
"Check Shanghai weather and add 1.0 plus 2.",
)],
test_tools(),
Some(true), // always enable thinking in this fixture
);
let expected_reasoning = "Need call the weather and add tools.";
let expected_text = "I will call the tools.";
@@ -261,9 +210,10 @@ async fn run_roundtrip_tool_call_mix(case: RoundtripCase) -> Result<()> {
AssistantContentBlock::ToolCall(AssistantToolCall {
id: "functions.add:1".to_string(),
name: "add".to_string(),
// Intentionally use a non-lexical order of keys and a different number
// formatting style to verify text-level fidelity of the roundtrip.
arguments: r#"{"y":1.00,"x":2}"#.to_string(),
// Intentionally use a non-lexical order of keys to verify text-level
// fidelity of the roundtrip where JSON formatting remains stable. The
// `items` key also exercises templates that call `arguments.items()`.
arguments: r#"{"y":1.0,"x":2,"items":["left","right"]}"#.to_string(),
}),
],
},
@@ -291,7 +241,7 @@ async fn run_roundtrip_tool_call_mix(case: RoundtripCase) -> Result<()> {
assert_eq!(tool_calls[1].name, "add");
assert_eq!(
tool_calls[1].arguments,
expected_arguments(&case, r#"{"y": 1.00, "x": 2}"#)?,
expected_arguments(&case, r#"{"y": 1.0, "x": 2, "items": ["left", "right"]}"#)?,
);
assert_eq!(
@@ -538,7 +488,6 @@ fn roundtrip_request(
request_id: impl Into<String>,
messages: Vec<ChatMessage>,
tools: Vec<ChatTool>,
thinking: Option<bool>,
) -> ChatRequest {
let mut request = ChatRequest {
request_id: request_id.into(),
@@ -552,12 +501,10 @@ fn roundtrip_request(
..ChatRequest::for_test()
};
// Explicitly enable or disable thinking so that rendering and parsing the reasoning block is
// exercised or skipped in the roundtrip. If unspecified, use the default template behavior.
if let Some(thinking) = thinking {
for key in ["thinking", "enable_thinking"] {
request.chat_options.template_kwargs.insert(key.to_string(), thinking.into());
}
// Enable thinking for some models so that rendering and parsing the reasoning block is
// exercised in the roundtrip.
for key in ["thinking", "enable_thinking"] {
request.chat_options.template_kwargs.insert(key.to_string(), true.into());
}
request
@@ -585,9 +532,13 @@ fn test_tools() -> Vec<ChatTool> {
"type": "object",
"properties": {
"y": { "type": "number" },
"x": { "type": "number" }
"x": { "type": "number" },
"items": {
"type": "array",
"items": { "type": "string" }
}
},
"required": ["y", "x"]
"required": ["y", "x", "items"]
}),
strict: None,
},
-11
View File
@@ -165,15 +165,6 @@ pub struct SharedRuntimeArgs {
#[serde(default)]
pub enable_log_requests: bool,
/// If specified, API server will add X-Request-Id header to responses.
#[arg(
long,
default_missing_value = "true",
num_args = 0..=1
)]
#[serde(default)]
pub enable_request_id_headers: bool,
/// Disable periodic logging of engine statistics (throughput, queue depth,
/// cache usage).
#[arg(long)]
@@ -247,7 +238,6 @@ impl SharedRuntimeArgs {
default_chat_template_kwargs: self.default_chat_template_kwargs,
chat_template_content_format: self.chat_template_content_format,
enable_log_requests: self.enable_log_requests,
enable_request_id_headers: self.enable_request_id_headers,
disable_log_stats: self.disable_log_stats,
grpc_port: self.grpc_port,
shutdown_timeout,
@@ -288,7 +278,6 @@ impl SharedRuntimeArgs {
default_chat_template_kwargs: self.default_chat_template_kwargs,
chat_template_content_format: self.chat_template_content_format,
enable_log_requests: self.enable_log_requests,
enable_request_id_headers: self.enable_request_id_headers,
disable_log_stats: self.disable_log_stats,
grpc_port: self.grpc_port,
shutdown_timeout,
-46
View File
@@ -43,7 +43,6 @@ fn serve_args_forward_python_flags_with_separator() {
default_chat_template_kwargs: None,
chat_template_content_format: Auto,
enable_log_requests: false,
enable_request_id_headers: false,
disable_log_stats: false,
served_model_name: [],
},
@@ -117,46 +116,6 @@ fn serve_args_accept_explicit_deepseek_v32_renderer() {
assert_eq!(args.runtime.renderer, RendererSelection::DeepSeekV32);
}
#[test]
fn serve_passes_enable_request_id_headers_into_config() {
let cli = Cli::try_parse_from([
"vllm-rs",
"serve",
"Qwen/Qwen3-0.6B",
"--enable-request-id-headers",
])
.unwrap();
let Command::Serve(args) = cli.command else {
panic!("expected serve args");
};
let config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string());
assert!(config.enable_request_id_headers);
}
#[test]
fn frontend_args_json_passes_enable_request_id_headers_into_config() {
let cli = Cli::try_parse_from([
"vllm-rs",
"frontend",
"--listen-fd",
"3",
"--input-address",
"ipc:///tmp/input.sock",
"--output-address",
"ipc:///tmp/output.sock",
"--args-json",
r#"{"model_tag":"Qwen/Qwen3-0.6B","enable_request_id_headers":true}"#,
])
.unwrap();
let Command::Frontend(args) = cli.command else {
panic!("expected frontend args");
};
let config = args.into_config();
assert!(config.enable_request_id_headers);
}
#[test]
fn serve_args_reject_unknown_renderer_value() {
let error = Cli::try_parse_from([
@@ -259,7 +218,6 @@ fn frontend_args_accept_json() {
default_chat_template_kwargs: None,
chat_template_content_format: Auto,
enable_log_requests: false,
enable_request_id_headers: false,
disable_log_stats: false,
served_model_name: [],
},
@@ -658,7 +616,6 @@ fn serve_args_accept_handshake_aliases() {
default_chat_template_kwargs: None,
chat_template_content_format: Auto,
enable_log_requests: false,
enable_request_id_headers: false,
disable_log_stats: false,
served_model_name: [],
},
@@ -776,7 +733,6 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() {
default_chat_template_kwargs: None,
chat_template_content_format: Auto,
enable_log_requests: false,
enable_request_id_headers: false,
disable_log_stats: false,
grpc_port: None,
shutdown_timeout: 0ns,
@@ -839,7 +795,6 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() {
default_chat_template_kwargs: None,
chat_template_content_format: Auto,
enable_log_requests: false,
enable_request_id_headers: false,
disable_log_stats: false,
grpc_port: None,
shutdown_timeout: 0ns,
@@ -917,7 +872,6 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present
default_chat_template_kwargs: None,
chat_template_content_format: Auto,
enable_log_requests: false,
enable_request_id_headers: false,
disable_log_stats: false,
grpc_port: None,
shutdown_timeout: 0ns,
+9
View File
@@ -620,6 +620,15 @@ pub struct ServerUnsupportedArgs {
#[arg(long)]
pub middleware: Option<Unsupported>,
/// If specified, API server will add X-Request-Id header to responses.
#[arg(
long,
visible_alias = "no-enable-request-id-headers",
default_missing_value = "true",
num_args = 0..=1
)]
pub enable_request_id_headers: Option<Unsupported>,
/// Disable FastAPI's OpenAPI schema, Swagger UI, and ReDoc endpoint.
#[arg(
long,
+2
View File
@@ -19,6 +19,7 @@ mod deepseek_r1;
mod delimited;
mod gemma4;
mod kimi;
mod minimax_m3;
mod qwen3;
use thiserror::Error;
@@ -29,6 +30,7 @@ pub use self::deepseek_r1::DeepSeekR1ReasoningParser;
pub(crate) use self::delimited::DelimitedReasoningParser;
pub use self::gemma4::Gemma4ReasoningParser;
pub use self::kimi::KimiReasoningParser;
pub use self::minimax_m3::MiniMaxM3ReasoningParser;
pub use self::qwen3::Qwen3ReasoningParser;
/// DeepSeek V3 currently shares the standard `<think>...</think>` parser.
@@ -0,0 +1,43 @@
use vllm_tokenizer::DynTokenizer;
use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningParser, Result};
/// Reasoning parser for MiniMax M3 style outputs.
///
/// MiniMax M3 uses `<mm:think>...</mm:think>` delimiters. Its chat template may
/// prefill either delimiter depending on the requested thinking mode, so the
/// shared delimited parser derives the starting state from the rendered prompt.
pub struct MiniMaxM3ReasoningParser {
inner: DelimitedReasoningParser,
}
impl MiniMaxM3ReasoningParser {
/// Create a MiniMax M3 parser backed by the shared delimited state machine.
pub fn new(tokenizer: DynTokenizer) -> Result<Self> {
Ok(Self {
inner: DelimitedReasoningParser::new(tokenizer, "<mm:think>", "</mm:think>", false)?,
})
}
}
impl ReasoningParser for MiniMaxM3ReasoningParser {
fn create(tokenizer: DynTokenizer) -> Result<Box<dyn ReasoningParser>>
where
Self: Sized + 'static,
{
Ok(Box::new(Self::new(tokenizer)?))
}
fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> {
self.inner.initialize(prompt_token_ids);
Ok(())
}
fn push(&mut self, delta: &str) -> Result<ReasoningDelta> {
Ok(self.inner.push(delta))
}
fn finish(&mut self) -> Result<ReasoningDelta> {
Ok(self.inner.finish())
}
}
+36 -1
View File
@@ -3,7 +3,8 @@ use std::sync::Arc;
use vllm_tokenizer::Tokenizer;
use super::{
DeepSeekR1ReasoningParser, DelimitedReasoningParser, Qwen3ReasoningParser, ReasoningParser,
DeepSeekR1ReasoningParser, DelimitedReasoningParser, MiniMaxM3ReasoningParser,
Qwen3ReasoningParser, ReasoningParser,
};
struct FakeTokenizer;
@@ -32,6 +33,8 @@ impl Tokenizer for FakeTokenizer {
"<|END_THINKING|>" => Some(4),
"◁think▷" => Some(5),
"◁/think▷" => Some(6),
"<mm:think>" => Some(8),
"</mm:think>" => Some(9),
_ => None,
}
}
@@ -159,3 +162,35 @@ fn deepseek_r1_stops_scanning_at_last_special_token() {
assert_eq!(delta.reasoning.as_deref(), Some("reason"));
assert_eq!(delta.content.as_deref(), Some("answer"));
}
#[test]
fn minimax_m3_handles_explicit_think_delimiters() {
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
let delta = parser.push("<mm:think>reason</mm:think>answer").unwrap();
assert_eq!(delta.reasoning.as_deref(), Some("reason"));
assert_eq!(delta.content.as_deref(), Some("answer"));
}
#[test]
fn minimax_m3_uses_prompt_prefilled_start_marker() {
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
parser.initialize(&[8]).unwrap();
let delta = parser.push("reason</mm:think>answer").unwrap();
assert_eq!(delta.reasoning.as_deref(), Some("reason"));
assert_eq!(delta.content.as_deref(), Some("answer"));
}
#[test]
fn minimax_m3_uses_prompt_prefilled_end_marker() {
let tokenizer = Arc::new(FakeTokenizer);
let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
parser.initialize(&[9]).unwrap();
let delta = parser.push("answer").unwrap();
assert_eq!(delta.reasoning, None);
assert_eq!(delta.content.as_deref(), Some("answer"));
}
@@ -68,7 +68,6 @@ async fn main() -> Result<()> {
default_chat_template_kwargs: None,
chat_template_content_format: ChatTemplateContentFormatOption::Auto,
enable_log_requests: false,
enable_request_id_headers: false,
disable_log_stats: false,
grpc_port: None,
shutdown_timeout: Duration::ZERO,
-2
View File
@@ -61,8 +61,6 @@ pub struct Config {
pub chat_template_content_format: ChatTemplateContentFormatOption,
/// Log a summary line for each completed request.
pub enable_log_requests: bool,
/// When `true`, set `X-Request-Id` on every HTTP response.
pub enable_request_id_headers: bool,
/// When `true`, suppress periodic stats logging (throughput, queue depth,
/// cache usage).
pub disable_log_stats: bool,
+1 -3
View File
@@ -85,9 +85,7 @@ async fn build_state(config: &Config) -> Result<Arc<AppState>> {
};
Ok(Arc::new(
AppState::new(served_model_names, chat)
.with_log_requests(config.enable_log_requests)
.with_request_id_headers(config.enable_request_id_headers),
AppState::new(served_model_names, chat).with_log_requests(config.enable_log_requests),
))
}
-2
View File
@@ -1,7 +1,5 @@
mod load;
mod metrics;
mod request_id;
pub use load::track_server_load;
pub use metrics::track_http_metrics;
pub use request_id::set_request_id_header;
@@ -1,24 +0,0 @@
use axum::extract::Request;
use axum::http::HeaderValue;
use axum::http::header::HeaderName;
use axum::middleware::Next;
use axum::response::Response;
use uuid::Uuid;
const X_REQUEST_ID: HeaderName = HeaderName::from_static("x-request-id");
/// Echo the request's `X-Request-Id` on the response, or generate a fresh
/// `uuid4` hex if the request did not provide one.
///
/// Original Python:
/// `vllm.entrypoints.openai.server_utils.XRequestIdMiddleware`.
pub async fn set_request_id_header(req: Request, next: Next) -> Response {
let incoming = req.headers().get(&X_REQUEST_ID).cloned();
let mut response = next.run(req).await;
let value = incoming.unwrap_or_else(|| {
HeaderValue::from_str(&Uuid::new_v4().simple().to_string())
.expect("uuid hex is valid header value")
});
response.headers_mut().insert(X_REQUEST_ID, value);
response
}
+2 -9
View File
@@ -56,18 +56,11 @@ fn build_router_with_dev_mode(state: Arc<AppState>, dev_mode_enabled: bool) -> R
.route("/is_sleeping", get(sleep::is_sleeping))
}
let enable_request_id_headers = state.enable_request_id_headers;
let mut router = router
router
.with_state(state.clone())
.layer(from_fn_with_state(state, middleware::track_server_load))
.layer(from_fn(middleware::track_http_metrics))
.layer(TraceLayer::new_for_http());
if enable_request_id_headers {
router = router.layer(from_fn(middleware::set_request_id_header));
}
router
.layer(TraceLayer::new_for_http())
}
#[cfg(test)]
@@ -3,33 +3,23 @@ mod types;
mod validate;
use std::collections::HashMap;
use std::convert::Infallible;
use std::result::Result;
use std::sync::Arc;
use asynk_strim_attr::{TryYielder, try_stream};
use axum::Json;
use axum::extract::State;
use axum::http::HeaderMap;
use axum::response::sse::{Event, Sse};
use axum::response::{IntoResponse, Response};
use futures::{Stream, StreamExt as _, pin_mut};
use thiserror_ext::AsReport as _;
use tracing::{error, info, trace};
use tracing::info;
use tracing_futures::Instrument as _;
use vllm_engine_core_client::protocol::logprobs::{Logprobs, PositionLogprobs};
use vllm_llm::{
CollectedGenerateOutput, FinishReason, GenerateOutput, GenerateOutputStreamExt as _,
};
use vllm_llm::{CollectedGenerateOutput, GenerateOutputStreamExt as _};
use self::convert::prepare_generate_request;
use self::types::{
GenerateLogprob, GenerateRequest, GenerateResponse, GenerateResponseChoice,
GenerateResponseStreamChoice, GenerateStreamResponse,
};
use crate::error::{ApiError, bail_server_error, server_error};
use self::types::{GenerateLogprob, GenerateRequest, GenerateResponse, GenerateResponseChoice};
use crate::error::{ApiError, server_error};
use crate::routes::openai::utils::logprobs::clamp_logprob;
use crate::routes::openai::utils::types::{ChatLogProbs, ChatLogProbsContent, TopLogProb, Usage};
use crate::routes::openai::utils::types::{ChatLogProbs, ChatLogProbsContent, TopLogProb};
use crate::routes::openai::utils::validated_json::ValidatedJson;
use crate::state::AppState;
use crate::utils::resolve_request_context;
@@ -56,7 +46,6 @@ pub async fn generate(
let log_request = state.enable_log_requests;
let include_logprobs = prepared.include_logprobs;
let include_prompt_logprobs = prepared.include_prompt_logprobs;
let stream = prepared.stream;
let raw_stream = match state
.chat
@@ -75,20 +64,6 @@ pub async fn generate(
}
};
if stream {
let chunk_stream = generate_chunk_stream(
raw_stream,
prepared.request_id,
log_request,
prepared.include_usage,
prepared.include_continuous_usage,
include_logprobs,
);
let sse_stream = generate_sse_stream(chunk_stream).instrument(request_span);
return Sse::new(sse_stream).into_response();
}
let collected = match raw_stream.collect_output().instrument(request_span.clone()).await {
Ok(collected) => collected,
Err(error) => {
@@ -123,102 +98,6 @@ pub async fn generate(
Json(response).into_response()
}
#[try_stream]
async fn generate_chunk_stream(
stream: impl Stream<Item = vllm_llm::Result<GenerateOutput>>,
request_id: String,
log_request: bool,
include_usage: bool,
include_continuous_usage: bool,
include_logprobs: bool,
mut y: TryYielder<GenerateStreamResponse, ApiError>,
) -> Result<(), ApiError> {
pin_mut!(stream);
let mut prompt_tokens: Option<u32> = None;
let mut output_tokens = 0_u32;
while let Some(next) = stream.next().await {
match next {
Ok(output) => {
if prompt_tokens.is_none() {
prompt_tokens =
output.prompt_info.as_ref().map(|info| info.prompt_token_ids.len() as u32);
}
let usage_prompt_tokens = prompt_tokens.unwrap_or_default();
let token_ids = output.token_ids;
output_tokens = output_tokens.saturating_add(token_ids.len() as u32);
let finish_reason = output.finish_reason;
if matches!(finish_reason.as_ref(), Some(FinishReason::Error)) {
bail_server_error!("Internal server error");
}
if let Some(finish_reason) = finish_reason.as_ref()
&& log_request
{
info!(
stream = true,
prompt_tokens = usage_prompt_tokens,
output_tokens,
finish_reason = finish_reason.as_str(),
"generate finished"
);
}
if token_ids.is_empty() && finish_reason.is_none() {
continue;
}
let logprobs = if include_logprobs && !token_ids.is_empty() {
let logprobs = output.logprobs.as_ref().ok_or_else(|| {
server_error!(
"raw generate stream requested logprobs but generation returned none"
)
})?;
Some(raw_logprobs_to_openai_chat(logprobs)?)
} else {
None
};
y.yield_ok(GenerateStreamResponse {
request_id: request_id.clone(),
choices: vec![GenerateResponseStreamChoice {
index: 0,
logprobs,
finish_reason: finish_reason.map(|reason| reason.as_str().to_string()),
token_ids,
}],
usage: include_continuous_usage
.then(|| Usage::from_counts(usage_prompt_tokens, output_tokens)),
})
.await;
}
Err(error) => {
error!(
error = %error.as_report(),
"raw generate stream failed"
);
bail_server_error!("{}", error.to_report_string());
}
}
}
if include_usage {
y.yield_ok(GenerateStreamResponse {
request_id,
choices: Vec::new(),
usage: Some(Usage::from_counts(
prompt_tokens.unwrap_or_default(),
output_tokens,
)),
})
.await;
}
Ok(())
}
fn collect_generate(
collected: CollectedGenerateOutput,
request_id: String,
@@ -334,94 +213,3 @@ fn position_to_logprob_map(position: &PositionLogprobs) -> HashMap<u32, Generate
fn format_token_id(token_id: u32) -> String {
format!("token_id:{token_id}")
}
/// Convert one raw-generate chunk stream into SSE events.
#[try_stream]
async fn generate_sse_stream(
stream: impl Stream<Item = Result<GenerateStreamResponse, ApiError>>,
mut y: TryYielder<Event, Infallible>,
) -> Result<(), Infallible> {
pin_mut!(stream);
while let Some(next) = stream.next().await {
match next {
Ok(chunk) => y.yield_ok(to_sse_event(&chunk)).await,
Err(error) => {
y.yield_ok(to_error_sse_event(&error)).await;
break;
}
}
}
y.yield_ok(done_sse_event()).await;
Ok(())
}
fn to_sse_event(chunk: &GenerateStreamResponse) -> Event {
let payload = serde_json::to_string(chunk).expect("generate chunk must serialize to JSON");
trace!(payload, "generate emitting chunk");
Event::default().data(payload)
}
fn to_error_sse_event(error: &ApiError) -> Event {
let payload = serde_json::to_string(&error.to_error_response())
.expect("ErrorResponse must serialize to JSON");
trace!(payload, "generate emitting error");
Event::default().data(payload)
}
fn done_sse_event() -> Event {
trace!("generate emitting done");
Event::default().data("[DONE]")
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use futures::{TryStreamExt as _, stream};
use vllm_llm::GeneratePromptInfo;
use super::*;
#[tokio::test]
async fn generate_chunk_stream_captures_late_prompt_info() {
let stream = stream::iter(vec![
Ok(GenerateOutput {
request_id: String::new(),
prompt_info: None,
token_ids: Vec::new(),
logprobs: None,
finish_reason: None,
kv_transfer_params: None,
}),
Ok(GenerateOutput {
request_id: String::new(),
prompt_info: Some(GeneratePromptInfo {
prompt_token_ids: Arc::from([11_u32, 22_u32]),
prompt_logprobs: None,
}),
token_ids: vec![33],
logprobs: None,
finish_reason: Some(FinishReason::stop_eos()),
kv_transfer_params: None,
}),
]);
let chunks: Vec<_> =
generate_chunk_stream(stream, "raw-stream".to_string(), false, true, true, false)
.try_collect()
.await
.expect("collect chunks");
assert_eq!(chunks.len(), 2);
assert_eq!(
chunks[0].usage.as_ref().expect("chunk usage").prompt_tokens,
2
);
assert_eq!(
chunks[1].usage.as_ref().expect("final usage").prompt_tokens,
2
);
}
}
@@ -10,9 +10,6 @@ use crate::utils::{ResolvedRequestContext, merge_kv_transfer_params};
pub struct PreparedRequest {
pub request_id: String,
pub text_request: TextRequest,
pub stream: bool,
pub include_usage: bool,
pub include_continuous_usage: bool,
pub include_logprobs: bool,
pub include_prompt_logprobs: bool,
}
@@ -26,18 +23,6 @@ pub fn prepare_generate_request(
) -> Result<PreparedRequest, ApiError> {
validate::validate_request_compat(&request, served_model_names)?;
let stream = request.stream;
let include_usage = request
.stream_options
.as_ref()
.and_then(|options| options.include_usage)
.unwrap_or(false);
let include_continuous_usage = include_usage
&& request
.stream_options
.as_ref()
.and_then(|options| options.continuous_usage_stats)
.unwrap_or(false);
let include_logprobs = request.sampling_params.logprobs.is_some();
let include_prompt_logprobs = request.sampling_params.prompt_logprobs.is_some();
let mut sampling_params = request.sampling_params;
@@ -62,9 +47,6 @@ pub fn prepare_generate_request(
Ok(PreparedRequest {
request_id: ctx.request_id,
text_request,
stream,
include_usage,
include_continuous_usage,
include_logprobs,
include_prompt_logprobs,
})
@@ -127,28 +109,4 @@ mod tests {
Some(json!({"connector": "x"}))
);
}
#[test]
fn prepare_generate_request_gates_continuous_usage_on_include_usage() {
let request: GenerateRequest = serde_json::from_value(json!({
"model": "Qwen/Qwen1.5-0.5B-Chat",
"token_ids": [11, 22],
"stream": true,
"stream_options": {
"continuous_usage_stats": true
},
"sampling_params": {}
}))
.expect("parse request");
let prepared = prepare_generate_request(
request,
&["Qwen/Qwen1.5-0.5B-Chat".to_string()],
ResolvedRequestContext::default(),
)
.expect("prepare");
assert!(!prepared.include_usage);
assert!(!prepared.include_continuous_usage);
}
}
@@ -5,7 +5,7 @@ use serde_json::{Map, Value};
use validator::Validate;
use vllm_text::SamplingParams;
use crate::routes::openai::utils::types::{ChatLogProbs, Normalizable, StreamOptions, Usage};
use crate::routes::openai::utils::types::{ChatLogProbs, Normalizable};
/// vLLM-compatible request type for the token-in/token-out generate API.
#[serde_with::skip_serializing_none]
@@ -17,7 +17,6 @@ pub struct GenerateRequest {
pub sampling_params: SamplingParams,
#[serde(default)]
pub stream: bool,
pub stream_options: Option<StreamOptions>,
pub cache_salt: Option<String>,
#[serde(default)]
pub priority: i32,
@@ -38,25 +37,6 @@ pub(super) struct GenerateResponseChoice {
pub token_ids: Vec<u32>,
}
/// Mirrors the Python vLLM `GenerateResponseStreamChoice` class.
#[serde_with::skip_serializing_none]
#[derive(Debug, Clone, Serialize)]
pub(super) struct GenerateResponseStreamChoice {
pub index: u32,
pub logprobs: Option<ChatLogProbs>,
pub finish_reason: Option<String>,
pub token_ids: Vec<u32>,
}
/// Mirrors the Python vLLM `GenerateStreamResponse` class.
#[serde_with::skip_serializing_none]
#[derive(Debug, Clone, Serialize)]
pub(super) struct GenerateStreamResponse {
pub request_id: String,
pub choices: Vec<GenerateResponseStreamChoice>,
pub usage: Option<Usage>,
}
/// Mirrors the Python vLLM `GenerateResponse` class.
#[serde_with::skip_serializing_none]
#[derive(Debug, Clone, Serialize)]
@@ -13,11 +13,8 @@ pub(super) fn validate_request_compat(
return Err(ApiError::model_not_found(model.clone()));
}
if request.stream_options.is_some() && !request.stream {
bail_invalid_request!(
param = "stream_options",
"stream_options are only supported when stream=true."
);
if request.stream {
bail_invalid_request!(param = "stream", "stream=true is not supported.");
}
if request.token_ids.is_empty() {
@@ -68,24 +65,11 @@ mod tests {
}
#[test]
fn validate_request_compat_accepts_streaming() {
fn validate_request_compat_rejects_streaming() {
let request = GenerateRequest {
stream: true,
..base_request()
};
assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_ok());
}
#[test]
fn validate_request_compat_rejects_stream_options_without_streaming() {
let request: GenerateRequest = serde_json::from_value(json!({
"model": "Qwen/Qwen1.5-0.5B-Chat",
"token_ids": [11, 22],
"stream": false,
"stream_options": {"include_usage": true},
"sampling_params": {}
}))
.expect("parse request");
assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_err());
}
+6 -327
View File
@@ -746,20 +746,6 @@ async fn test_app() -> axum::Router {
)))
}
async fn test_app_with_request_id_headers() -> (axum::Router, MockEngineTask) {
let (chat, engine_task) = test_models_with_engine_outputs_and_backend(
b"engine-openai-request-id",
default_stream_output_specs(),
Arc::new(FakeChatBackend::new()),
)
.await;
let app = build_router(Arc::new(
AppState::new(vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()], chat)
.with_request_id_headers(true),
));
(app, engine_task)
}
async fn test_health_app_with_engine_script<F>(
script: F,
) -> (axum::Router, Arc<AppState>, MockEngineTask)
@@ -961,18 +947,6 @@ async fn health_status(app: &axum::Router) -> (StatusCode, Bytes) {
(status, body)
}
async fn health_response(app: &axum::Router, request_id: Option<&str>) -> axum::response::Response {
let mut builder = Request::builder().method("GET").uri("/health");
if let Some(request_id) = request_id {
builder = builder.header("X-Request-Id", request_id);
}
app.clone()
.call(builder.body(Body::empty()).expect("build request"))
.await
.expect("call app")
}
fn metric_value(rendered: &str, metric: &str, labels: Option<&str>) -> Option<f64> {
rendered.lines().find_map(|line| {
let rest = line.strip_prefix(metric)?;
@@ -1020,43 +994,6 @@ async fn list_models_returns_configured_model() {
assert_eq!(json["data"][0]["id"], "Qwen/Qwen1.5-0.5B-Chat");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn request_id_header_is_absent_by_default() {
let app = test_app().await;
let response = health_response(&app, None).await;
assert_eq!(response.status(), StatusCode::OK);
assert!(!response.headers().contains_key("x-request-id"));
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn request_id_header_generates_uuid_hex_when_enabled() {
let (app, _engine_task) = test_app_with_request_id_headers().await;
let response = health_response(&app, None).await;
assert_eq!(response.status(), StatusCode::OK);
let request_id = response
.headers()
.get("x-request-id")
.expect("x-request-id header")
.to_str()
.expect("header is ascii");
assert_eq!(request_id.len(), 32);
assert!(request_id.chars().all(|ch| ch.is_ascii_hexdigit()));
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn request_id_header_echoes_incoming_header_when_enabled() {
let (app, _engine_task) = test_app_with_request_id_headers().await;
let response = health_response(&app, Some("req-123")).await;
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(response.headers().get("x-request-id").unwrap(), "req-123");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn version_returns_engine_vllm_version() {
@@ -2488,72 +2425,8 @@ async fn non_stream_raw_generate_returns_token_output_envelope() {
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn stream_raw_generate_returns_sse_chunks_and_usage() {
let ipc = IpcNamespace::new().expect("create ipc namespace");
let handshake_address = ipc.handshake_endpoint();
let engine_id = b"engine-raw-generate-stream".to_vec();
let engine_task = MockEngineTask::new(spawn_mock_engine_task(
handshake_address.clone(),
engine_id.clone(),
|dealer, push| {
boxed_test_future(async move {
let add = recv_engine_message(dealer).await;
let request: EngineCoreRequest =
rmp_serde::from_slice(&add[1]).expect("decode request");
assert_eq!(request.prompt_token_ids.as_deref(), Some(&[11, 22][..]));
assert_eq!(request.external_req_id.as_deref(), Some("raw-stream"));
send_outputs(
push,
EngineCoreOutputs {
engine_index: 0,
outputs: vec![
request_output_with_logprobs(
&request.request_id,
vec![33],
None,
None,
Some(sample_logprobs_for_token(33, 34)),
None,
),
request_output_with_logprobs(
&request.request_id,
vec![44],
Some(EngineCoreFinishReason::Stop),
None,
Some(sample_logprobs_for_token(44, 45)),
None,
),
],
scheduler_stats: None,
timestamp: 0.0,
utility_output: None,
finished_requests: None,
wave_complete: None,
start_wave: None,
},
)
.await;
})
},
));
let client = EngineCoreClient::connect(
EngineCoreClientConfig::new_single(handshake_address)
.with_model_name("test-model")
.with_local_input_output_addresses(
Some(ipc.input_endpoint()),
Some(ipc.output_endpoint()),
),
)
.await
.expect("connect client");
let chat = ChatLlm::from_shared_backend(Llm::new(client), Arc::new(FakeChatBackend::new()));
let mut app = build_router(Arc::new(AppState::new(
vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()],
chat,
)));
async fn raw_generate_rejects_streaming() {
let mut app = test_app().await;
let response = app
.call(
@@ -2564,17 +2437,9 @@ async fn stream_raw_generate_returns_sse_chunks_and_usage() {
.body(Body::from(
json!({
"model": "Qwen/Qwen1.5-0.5B-Chat",
"request_id": "raw-stream",
"token_ids": [11, 22],
"stream": true,
"stream_options": {
"include_usage": true,
"continuous_usage_stats": true
},
"sampling_params": {
"max_tokens": 2,
"logprobs": 1
}
"sampling_params": {}
})
.to_string(),
))
@@ -2583,196 +2448,10 @@ async fn stream_raw_generate_returns_sse_chunks_and_usage() {
.await
.expect("call app");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response.headers().get("content-type").and_then(|value| value.to_str().ok()),
Some("text/event-stream")
);
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
engine_task.await.expect("mock engine task");
let text = String::from_utf8(body.to_vec()).expect("utf8 body");
let payloads = sse_data_payloads(&text);
assert_eq!(payloads.len(), 4, "{text}");
let first: serde_json::Value = serde_json::from_str(payloads[0]).expect("first chunk json");
assert_eq!(first["request_id"], "raw-stream");
assert_eq!(first["choices"][0]["index"], 0);
assert_eq!(first["choices"][0]["token_ids"], json!([33]));
assert_eq!(
first["choices"][0]["logprobs"]["content"][0]["token"],
"token_id:33"
);
assert_eq!(first["usage"]["prompt_tokens"], 2);
assert_eq!(first["usage"]["completion_tokens"], 1);
let second: serde_json::Value = serde_json::from_str(payloads[1]).expect("second chunk json");
assert_eq!(second["choices"][0]["token_ids"], json!([44]));
assert_eq!(second["choices"][0]["finish_reason"], "stop");
assert_eq!(second["usage"]["completion_tokens"], 2);
let usage: serde_json::Value = serde_json::from_str(payloads[2]).expect("usage chunk json");
assert_eq!(usage["choices"], json!([]));
assert_eq!(usage["usage"]["prompt_tokens"], 2);
assert_eq!(usage["usage"]["completion_tokens"], 2);
assert_eq!(usage["usage"]["total_tokens"], 4);
assert_eq!(payloads[3], "[DONE]");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn stream_raw_generate_emits_final_usage_without_continuous_usage() {
let (mut app, engine_task) = test_app_with_stream_output_specs(vec![
(vec![33], None),
(vec![44], Some(EngineCoreFinishReason::Stop)),
])
.await;
let response = app
.call(
Request::builder()
.method("POST")
.uri("/inference/v1/generate")
.header("content-type", "application/json")
.body(Body::from(
json!({
"model": "Qwen/Qwen1.5-0.5B-Chat",
"request_id": "raw-stream-final-usage",
"token_ids": [11, 22],
"stream": true,
"stream_options": {
"include_usage": true
},
"sampling_params": {
"max_tokens": 2
}
})
.to_string(),
))
.expect("build request"),
)
.await
.expect("call app");
assert_eq!(response.status(), StatusCode::OK);
let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
engine_task.await.expect("mock engine task");
let text = String::from_utf8(body.to_vec()).expect("utf8 body");
let payloads = sse_data_payloads(&text);
assert_eq!(payloads.len(), 4, "{text}");
let first: serde_json::Value = serde_json::from_str(payloads[0]).expect("first chunk json");
assert_eq!(first["choices"][0]["token_ids"], json!([33]));
assert!(first.get("usage").is_none());
let second: serde_json::Value = serde_json::from_str(payloads[1]).expect("second chunk json");
assert_eq!(second["choices"][0]["token_ids"], json!([44]));
assert_eq!(second["choices"][0]["finish_reason"], "stop");
assert!(second.get("usage").is_none());
let usage: serde_json::Value = serde_json::from_str(payloads[2]).expect("usage chunk json");
assert_eq!(usage["choices"], json!([]));
assert_eq!(usage["usage"]["prompt_tokens"], 2);
assert_eq!(usage["usage"]["completion_tokens"], 2);
assert_eq!(usage["usage"]["total_tokens"], 4);
assert_eq!(payloads[3], "[DONE]");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn stream_raw_generate_emits_empty_finish_chunk() {
let (mut app, engine_task) = test_app_with_stream_output_specs(vec![
(vec![33], None),
(vec![], Some(EngineCoreFinishReason::Stop)),
])
.await;
let response = app
.call(
Request::builder()
.method("POST")
.uri("/inference/v1/generate")
.header("content-type", "application/json")
.body(Body::from(
json!({
"model": "Qwen/Qwen1.5-0.5B-Chat",
"request_id": "raw-stream-empty-finish",
"token_ids": [11, 22],
"stream": true,
"sampling_params": {
"max_tokens": 2
}
})
.to_string(),
))
.expect("build request"),
)
.await
.expect("call app");
assert_eq!(response.status(), StatusCode::OK);
let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
engine_task.await.expect("mock engine task");
let text = String::from_utf8(body.to_vec()).expect("utf8 body");
let payloads = sse_data_payloads(&text);
assert_eq!(payloads.len(), 3, "{text}");
let first: serde_json::Value = serde_json::from_str(payloads[0]).expect("first chunk json");
assert_eq!(first["choices"][0]["token_ids"], json!([33]));
assert!(first["choices"][0].get("finish_reason").is_none());
let second: serde_json::Value = serde_json::from_str(payloads[1]).expect("second chunk json");
assert_eq!(second["choices"][0]["token_ids"], json!([]));
assert_eq!(second["choices"][0]["finish_reason"], "stop");
assert_eq!(payloads[2], "[DONE]");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn stream_raw_generate_error_finish_returns_sse_error() {
let (mut app, engine_task) =
test_app_with_stream_output_specs(vec![(vec![], Some(EngineCoreFinishReason::Error))])
.await;
let response = app
.call(
Request::builder()
.method("POST")
.uri("/inference/v1/generate")
.header("content-type", "application/json")
.body(Body::from(
json!({
"model": "Qwen/Qwen1.5-0.5B-Chat",
"request_id": "raw-stream-error",
"token_ids": [11, 22],
"stream": true,
"stream_options": {
"include_usage": true
},
"sampling_params": {
"max_tokens": 2
}
})
.to_string(),
))
.expect("build request"),
)
.await
.expect("call app");
assert_eq!(response.status(), StatusCode::OK);
let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
engine_task.await.expect("mock engine task");
let text = String::from_utf8(body.to_vec()).expect("utf8 body");
assert!(text.contains("\"type\":\"server_error\""), "{text}");
assert!(text.contains("Internal server error"), "{text}");
assert!(!text.contains("\"finish_reason\":\"error\""), "{text}");
assert!(!text.contains("\"usage\":"), "{text}");
assert!(text.trim_end().ends_with("data: [DONE]"), "{text}");
let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
assert_eq!(json["error"]["param"], "stream");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
-9
View File
@@ -17,8 +17,6 @@ pub struct AppState {
pub chat: ChatLlm,
/// Whether to log a summary line for each completed request.
pub enable_log_requests: bool,
/// Whether to set X-Request-Id on every HTTP response.
pub enable_request_id_headers: bool,
/// Number of in-flight inference requests currently owned by this frontend.
server_load: AtomicU64,
}
@@ -41,7 +39,6 @@ impl AppState {
served_model_names,
chat,
enable_log_requests: false,
enable_request_id_headers: false,
server_load: AtomicU64::new(0),
}
}
@@ -52,12 +49,6 @@ impl AppState {
self
}
/// Enable X-Request-Id response headers.
pub fn with_request_id_headers(mut self, enabled: bool) -> Self {
self.enable_request_id_headers = enabled;
self
}
/// The primary model name echoed back in API responses (the first served
/// name).
pub fn primary_model_name(&self) -> &str {
+23
View File
@@ -0,0 +1,23 @@
[package]
name = "vllm-tool-parser-py"
version.workspace = true
edition.workspace = true
license.workspace = true
[lib]
name = "_rust_tool_parser"
crate-type = ["cdylib", "rlib"]
[features]
default = []
extension-module = ["pyo3/extension-module"]
[dependencies]
pyo3.workspace = true
pythonize = { workspace = true, features = ["serde_json"] }
serde_json.workspace = true
thiserror-ext.workspace = true
vllm-tool-parser.workspace = true
[lints]
workspace = true
+367
View File
@@ -0,0 +1,367 @@
//! Thin PyO3 bindings for `vllm_tool_parser`.
//!
//! This crate exposes the Rust tool parser trait and data shapes to Python
//! while keeping parser state, grammar, and schema-aware argument conversion in
//! Rust. Python callers should use this module as a typed bridge and keep any
//! vLLM protocol adaptation outside the binding.
use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use pyo3::types::{PyAny, PyModule};
use pythonize::{depythonize, pythonize};
use serde_json::Value;
use thiserror_ext::AsReport as _;
use vllm_tool_parser::{Tool, ToolCallDelta, ToolParser, ToolParserOutput};
macro_rules! tool_parser_factory {
($($parser:ident),+ $(,)?) => {
fn create_tool_parser(
name: &str,
tools: &[Tool],
) -> PyResult<Box<dyn ToolParser>> {
match name {
$(
stringify!($parser) => {
<vllm_tool_parser::$parser as ToolParser>::create(tools)
}
)+
_ => {
return Err(PyValueError::new_err(format!(
"unsupported tool parser `{name}`"
)));
}
}
.map_err(|error| PyValueError::new_err(error.to_report_string()))
}
};
}
// Export a tool parser to Python by registering it here.
tool_parser_factory! {
DeepSeekV4ToolParser,
MinimaxM3ToolParser,
}
#[pyclass(name = "Tool", module = "vllm._rust_tool_parser", skip_from_py_object)]
#[derive(Clone)]
struct PyTool(Tool);
#[pymethods]
impl PyTool {
#[new]
#[pyo3(signature = (name, description, parameters, strict=None))]
fn new(
name: String,
description: Option<String>,
parameters: &Bound<'_, PyAny>,
strict: Option<bool>,
) -> PyResult<Self> {
let parameters = depythonize::<Value>(parameters).map_err(|error| {
PyValueError::new_err(format!(
"failed to convert tool parameters from Python to JSON: {error}"
))
})?;
Ok(Self(Tool {
name,
description,
parameters,
strict,
}))
}
#[getter]
fn name(&self) -> &str {
&self.0.name
}
#[getter]
fn description(&self) -> Option<&str> {
self.0.description.as_deref()
}
#[getter]
fn parameters(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
pythonize(py, &self.0.parameters).map(Bound::unbind).map_err(|error| {
PyValueError::new_err(format!(
"failed to convert tool parameters from JSON to Python: {error}"
))
})
}
#[getter]
fn strict(&self) -> Option<bool> {
self.0.strict
}
}
#[pyclass(
name = "ToolCallDelta",
module = "vllm._rust_tool_parser",
skip_from_py_object
)]
#[derive(Clone)]
struct PyToolCallDelta(ToolCallDelta);
#[pymethods]
impl PyToolCallDelta {
#[new]
#[pyo3(signature = (tool_index, name, arguments))]
fn new(tool_index: usize, name: Option<String>, arguments: String) -> Self {
Self(ToolCallDelta {
tool_index,
name,
arguments,
})
}
#[getter]
fn tool_index(&self) -> usize {
self.0.tool_index
}
#[getter]
fn name(&self) -> Option<&str> {
self.0.name.as_deref()
}
#[getter]
fn arguments(&self) -> &str {
&self.0.arguments
}
}
#[pyclass(
name = "ToolParserOutput",
module = "vllm._rust_tool_parser",
skip_from_py_object
)]
#[derive(Clone)]
struct PyToolParserOutput(ToolParserOutput);
#[pymethods]
impl PyToolParserOutput {
#[new]
#[pyo3(signature = (normal_text="", calls=None))]
fn new(py: Python<'_>, normal_text: &str, calls: Option<Vec<Py<PyToolCallDelta>>>) -> Self {
let calls =
calls.unwrap_or_default().iter().map(|call| call.borrow(py).0.clone()).collect();
Self(ToolParserOutput {
normal_text: normal_text.to_owned(),
calls,
})
}
#[getter]
fn normal_text(&self) -> &str {
&self.0.normal_text
}
#[getter]
fn calls(&self) -> Vec<PyToolCallDelta> {
self.0.calls.iter().cloned().map(PyToolCallDelta).collect()
}
fn append(&mut self, other: PyRef<'_, PyToolParserOutput>) {
self.0.append(other.0.clone());
}
fn coalesce_calls(&self) -> Self {
Self(self.0.clone().coalesce_calls())
}
}
#[pyclass(name = "ToolParser", module = "vllm._rust_tool_parser", unsendable)]
struct PyToolParser(Box<dyn ToolParser>);
impl PyToolParser {
fn parse_into_output(&mut self, chunk: &str, output: &mut PyToolParserOutput) -> PyResult<()> {
self.0
.parse_into(chunk, &mut output.0)
.map_err(|error| PyValueError::new_err(error.to_report_string()))
}
}
#[pymethods]
impl PyToolParser {
#[new]
fn new(py: Python<'_>, parser_name: &str, tools: Vec<Py<PyTool>>) -> PyResult<Self> {
let tools = tools.iter().map(|tool| tool.borrow(py).0.clone()).collect::<Vec<_>>();
create_tool_parser(parser_name, &tools).map(Self)
}
fn parse_into(
&mut self,
chunk: &str,
mut output: PyRefMut<'_, PyToolParserOutput>,
) -> PyResult<()> {
self.parse_into_output(chunk, &mut output)
}
fn finish(&mut self) -> PyResult<PyToolParserOutput> {
self.0
.finish()
.map(PyToolParserOutput)
.map_err(|error| PyValueError::new_err(error.to_report_string()))
}
fn reset(&mut self) -> String {
self.0.reset()
}
fn preserve_special_tokens(&self) -> bool {
self.0.preserve_special_tokens()
}
}
#[pymodule]
fn _rust_tool_parser(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<PyTool>()?;
m.add_class::<PyToolCallDelta>()?;
m.add_class::<PyToolParserOutput>()?;
m.add_class::<PyToolParser>()?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
fn with_python<R>(f: impl for<'py> FnOnce(Python<'py>) -> R) -> R {
Python::initialize();
Python::attach(f)
}
fn tool_schema() -> Value {
json!({
"type": "object",
"properties": {
"user_id": {"type": "integer"},
"shipping": {
"type": "object",
"properties": {
"city": {"type": "string"},
"zip": {"type": "integer"}
}
}
}
})
}
fn build_call() -> String {
r#"<DSMLtool_calls>
<DSMLinvoke name="create_order">
<DSMLparameter name="user_id" string="false">42</DSMLparameter>
<DSMLparameter name="shipping" string="false">{"city":"Singapore","zip":18956}</DSMLparameter>
</DSMLinvoke>
</DSMLtool_calls>"#
.to_owned()
}
fn make_py_tool(py: Python<'_>) -> PyResult<Py<PyTool>> {
let parameters = pythonize(py, &tool_schema()).map_err(|error| {
PyValueError::new_err(format!(
"failed to convert test schema from JSON to Python: {error}"
))
})?;
Py::new(
py,
PyTool::new(
"create_order".to_owned(),
Some("Create an order".to_owned()),
&parameters,
None,
)?,
)
}
#[test]
fn tool_round_trips_typed_fields() {
with_python(|py| {
let tool = make_py_tool(py)?;
let borrowed = tool.borrow(py);
assert_eq!(borrowed.name(), "create_order");
assert_eq!(borrowed.description(), Some("Create an order"));
assert_eq!(borrowed.strict(), None);
let parameters = borrowed.parameters(py)?;
let parameters = depythonize::<Value>(parameters.bind(py))?;
assert_eq!(parameters, tool_schema());
PyResult::Ok(())
})
.unwrap();
}
#[test]
fn output_append_and_coalesce_calls() {
with_python(|py| {
let first = Py::new(
py,
PyToolCallDelta::new(0, Some("create_order".to_owned()), "{\"a\"".to_owned()),
)?;
let second = Py::new(py, PyToolCallDelta::new(0, None, ":1}".to_owned()))?;
let mut output = PyToolParserOutput::new(py, "text", Some(vec![first]));
let other = Py::new(py, PyToolParserOutput::new(py, "", Some(vec![second])))?;
output.append(other.borrow(py));
let coalesced = output.coalesce_calls();
assert_eq!(coalesced.normal_text(), "text");
let calls = coalesced.calls();
assert_eq!(calls.len(), 1);
assert_eq!(calls[0].tool_index(), 0);
assert_eq!(calls[0].name(), Some("create_order"));
assert_eq!(calls[0].arguments(), "{\"a\":1}");
PyResult::Ok(())
})
.unwrap();
}
#[test]
fn parser_parse_finish_and_preserve_special_tokens() {
with_python(|py| {
let tool = make_py_tool(py)?;
let mut parser = PyToolParser::new(py, "DeepSeekV4ToolParser", vec![tool])?;
assert!(parser.preserve_special_tokens());
let mut output = PyToolParserOutput::new(py, "", None);
parser.parse_into_output(&build_call(), &mut output)?;
let finish = Py::new(py, parser.finish()?)?;
output.append(finish.borrow(py));
let output = output.coalesce_calls();
assert_eq!(output.normal_text(), "");
let calls = output.calls();
assert_eq!(calls.len(), 1);
assert_eq!(calls[0].name(), Some("create_order"));
assert_eq!(
serde_json::from_str::<Value>(calls[0].arguments()).unwrap(),
json!({
"user_id": 42,
"shipping": {
"city": "Singapore",
"zip": 18956
}
})
);
assert_eq!(parser.reset(), "");
PyResult::Ok(())
})
.unwrap();
}
#[test]
fn parser_errors_for_unknown_name() {
with_python(|py| {
let tool = make_py_tool(py)?;
let error = match PyToolParser::new(py, "missing", vec![tool]) {
Ok(_) => panic!("missing parser name unexpectedly succeeded"),
Err(error) => error,
};
let message = format!("{error}");
assert!(message.contains("unsupported tool parser `missing`"));
PyResult::Ok(())
})
.unwrap();
}
}
+1 -1
View File
@@ -8,7 +8,7 @@ const HERMES_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
marker_whitespace: JsonToolCallWhitespace::Optional,
delimiter: None,
name_key: "name",
arguments_key: &["arguments"],
arguments_key: "arguments",
};
/// Tool parser for Hermes XML-wrapped JSON tool calls.

Some files were not shown because too many files have changed in this diff Show More