Compare commits

...
Author SHA1 Message Date
Zijing LiuandYongye Zhu 8b00f41237 [MiniMax-M3] Retune decode indexer launch grid for occupancy (#130)
Signed-off-by: Zijing Liu <liuzijing2014@gmail.com>
2026-06-25 18:14:52 +00:00
Xin LiandGitHub 0586a20184 Add MiniMax-M3 modelopt nvfp4 support (#46380)
Signed-off-by: Xin Li <xinli@nvidia.com>
2026-06-25 11:11:41 -07:00
Yongye ZhuandClaude Opus 4.8 7a672233eb Bump flashinfer-jit-cache to 0.6.13rc2
Align flashinfer-jit-cache with flashinfer-python/flashinfer-cubin, which
are already pinned to 0.6.13rc2 in requirements/cuda.txt.

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

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
2026-06-23 00:01:30 +00:00
Yongye Zhu ea890c8066 remove contiguous constraint for prefill_topk
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
2026-06-22 05:21:19 +00:00
Thien TranandRoger Wang 8717eccaf4 Fix non-contiguous topk buffer for MSA sparse GQA prefill (#105)
(cherry picked from commit eff83fa5c1ecf47678a262ed85f5d751b5cf1b67)
2026-06-20 13:53:02 -07:00
Zijing LiuandRoger Wang 9ac6a8ed61 NVFP4 TRTLLM-Gen fused MoE: clamped SwiGLU-OAI support + NVFP4 (#97)
Signed-off-by: Zijing Liu <liuzijing2014@gmail.com>
(cherry picked from commit 20bafcab50f1f63de1e0a5f810d7b844888fc0ac)
2026-06-20 13:52:56 -07:00
Yongye ZhuandRoger Wang ef762c26e5 [MoE] Plumb gemm1_alpha/beta/clamp_limit into TRT-LLM FP8 MoE
The FlashInfer FP8 block-scale MoE kernels (trtllm_fp8_block_scale_moe
and trtllm_fp8_block_scale_routed_moe) accept optional per-expert SwiGLU
parameters gemm1_alpha, gemm1_beta, gemm1_clamp_limit that realize the
OAI SwiGLU variant for MXFP8. The FP8 experts did not pass them.

- Build per-expert alpha/beta/clamp tensors in TrtLlmFp8ExpertsBase and
  pass them to the block-scale (monolithic) and routed (modular) kernels,
  mirroring the existing MXFP4 experts.
- Forward gemm1_alpha/gemm1_beta from layer.swiglu_alpha/beta in the
  MXFP8-capable FP8 quant-config builders (fp8, online mxfp8,
  compressed-tensors mxfp8). Previously only swiglu_limit was forwarded;
  only ModelOptMxfp8 forwarded all three, so other paths silently dropped
  alpha/beta.
- Add swigluoai_uninterleave to the supported activations (maps to
  FlashInfer Swiglu; the OAI behavior comes from the gemm1_* params).

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
(cherry picked from commit f4b8540570)
2026-06-20 13:51:49 -07:00
Yongye ZhuandRoger Wang 985d00a751 [Test] M3 indexer: add sm_scale arg to _reference_index_topk
Fixes mypy call-arg errors: two fp8 call sites pass sm_scale (by
keyword and positionally) that the reference signature lacked. Add
sm_scale (default 1.0) and apply it to the score; top-k selection is
invariant to a positive scalar, so existing callers are unchanged.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
(cherry picked from commit 522f262265)
2026-06-20 13:51:43 -07:00
Yongye ZhuandRoger Wang 70ae0bb9e5 run indexer eager
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
(cherry picked from commit e6dd784466)
2026-06-20 13:51:43 -07:00
Thien TranandRoger Wang 18d9bd9399 [Build] fmha_sm100: install indexer csrc/cutlass + package_data globs
PR #93 integrated the MSA indexer but did not update the cmake install
rules / setup.py package_data to vendor the new files. Add install rules
for csrc/, cutlass/include, cutlass/tools/util/include and the matching
package_data globs.

(cherry picked from commit ea59ac42ab8e4a68f3bc362034de5199b59246df)

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
(cherry picked from commit e6fe1ecd64)
2026-06-20 13:51:43 -07:00
Yongye ZhuandRoger Wang 28ee7f57df [Model] M3 MSA indexer: align decode path with rebased index API
The MSA builder/impl were written against the pre-#45743 indexer decode
API. After rebasing onto main, MiniMaxM3IndexerDecodeMetadata gained a
required max_decode_query_len field and minimax_m3_index_decode dropped
its sm_scale parameter, so the MSA path crashed at engine init during
cudagraph profiling (missing max_decode_query_len).

- Pass max_decode_query_len when building decode metadata.
- Drop the stale sm_scale arg and pass max_decode_query_len in the decode
  kernel call, matching the Triton builder/impl.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
(cherry picked from commit 0b70ba38a1)
2026-06-20 13:51:43 -07:00
Yongye ZhuandRoger Wang c12ed89131 [Model] M3 MSA indexer: swap decode to Triton fused kernel, add fp8
Route the MSA (SM100) indexer decode through the Triton fused
minimax_m3_index_decode kernel (the same kernel the Triton indexer impl
uses) instead of fmha_sm100's OnlyScore path. For q_len==1 decode the
Triton kernel is a purpose-built vector x matrix score with a 256-way
split-K and fused split-K top-k, which beats fmha's OnlyScore (wasted MMA
tiles on a single query, 64-split cap) by ~1.1-3.7x in benchmarks. It is
cudagraph-safe by construction (shape-constant split grids) and writes the
shared topk_indices_buffer via out=. Prefill keeps fmha OnlyScore + the
single-pass Triton top-k, where fmha is ~3-5x faster for the wide score.

This drops the persistent fmha decode plan buffers, the num_kv_splits
estimate, and the now-unused VLLM_M3_INDEXER_CONTEXT_LEN env.

Add fp8 (e4m3) index KV cache support to the Triton decode score kernel:
make the QK MMA accumulate in fp32 (out_dtype=tl.float32) so the per-block
max score is exact for the e4m3 cache. Top-k is invariant to positive
scalar scaling, so fp8 needs no scale tuning.

Tests: tests/kernels/attention/test_minimax_m3.py (44 passed), including a
new test_decode_index_topk_fp8 validating fp8 decode top-k vs a
dequantized-fp32 reference. mypy + ruff clean. AI assistance was used.

Co-authored-by: Claude <noreply@anthropic.com>
Signed-off-by: Yongye Zhu <yongye@inferact.ai>

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
(cherry picked from commit 0a102417bb)
2026-06-20 13:51:43 -07:00
Yongye ZhuandRoger Wang 703e492aab [Model] M3: tune indexer decode split heuristic from fp8 sweep
estimate_num_kv_splits was a static replica of fmha_sm100's auto cost heuristic,
which over-splits the decode score: an fp8 e4m3 cudagraph-timed split sweep shows
the fixed-split (cudagraph-safe) path hits a sharp split-KV cliff above ~16
splits with >1 work row -- up to ~9x slower (e.g. reqs=2 ctx=65536: s64=153us vs
s16=15us). The previous formula picked 52-64 splits there.

Replace it with a measured, static fit (no kernel calls, so it stays
cudagraph-stable and sync-free):
- bs==1: min(64, kv_tiles) -- no cliff, fill with many splits.
- bs>=2: min(num_sms // work_rows, cap, kv_tiles), where the cliff-safe cap is
  16 through the ~60-100k context target and 32 for longer context (the cliff
  ceiling rises with context).

Across an fp8 cudagraph sweep this is within ~3% of the best fixed split for
every (batch, context) measured. The planner's own auto path is faster for
bs>=2 (adaptive split distribution) but is catastrophic for bs==1 (up to ~11x)
and needs a device->host sync, so it is unsuitable for the captured decode path.

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

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
(cherry picked from commit 6dd73260e8)
2026-06-20 13:51:43 -07:00
Yongye ZhuandRoger Wang 5f0bafab12 [Model] M3: read sparse-attend top-k from the persistent buffer
Fixes a cudagraph-replay AssertionError on bs=1 decode: with the indexer in the
captured segment and only the attend eager-broken, the (decode_topk, prefill_topk)
Python tuple handed across the eager break was frozen at capture and replayed
stale against the current step's metadata.

The attend now reads its top-k directly from the shared persistent
``topk_indices_buffer`` (decode at [:, :nd], prefill at [:, nd:num_tokens]),
sliced by the current step's metadata -- nothing crosses the eager break as a
Python value. The indexer (captured) writes the buffer; the attend (eager) reads
it; the breakable-cudagraph segment ordering guarantees the write precedes the
read on replay.

- MiniMaxM3SparseImpl.forward (+ Triton and MSA subclasses): drop the topk_idx
  arg, read layer.topk_indices_buffer.
- nvidia model: store the buffer on the sparse-attention layer; narrow the eager
  break so the indexer is captured and only the attend (_run_sparse_attn) is
  eager-broken.
- amd model: store the buffer on the layer; indexer + attend stay in one eager
  break (no capture), but the attend reads the buffer too.

42/42 in test_minimax_m3.py pass.

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

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
(cherry picked from commit 627f730369)
2026-06-20 13:51:43 -07:00
Yongye ZhuandRoger Wang 545393a3ab [Model] M3: log indexer kernel selection in select_indexer_impl_cls
Emit an info_once line naming the chosen indexer impl (MSA fmha_sm100 vs Triton)
and the deciding inputs (topk_blocks, indexer_kv_dtype, sm100), so the active
kernel path is visible at startup.

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

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
(cherry picked from commit 7fd090bd9c)
2026-06-20 13:51:43 -07:00
Yongye ZhuandRoger Wang d71445af45 [Model] M3: route Triton indexer top-k through shared topk_indices_buffer
Unify the Triton indexer impl onto the same persistent top-k buffer as the MSA
impl, and thread the buffer through the AMD model too.

- minimax_m3_index_decode gains an out= param (writes out[:, :total_q]); the merge
  kernel already writes via strides, so a buffer view works.
- MiniMaxM3IndexerTritonImpl.forward writes decode ([:, :nd]) and prefill ([:, nd:])
  into the shared topk_indices_buffer and returns views into it (no fresh per-step
  top-k allocations), matching the MSA impl.
- amd/model.py: allocate the model-level topk_indices_buffer and thread it
  model -> decoder layer -> sparse attention -> indexer (mirrors nvidia). AMD keeps
  its eager break, so this is purely allocation reuse there.

test_msa_indexer_impl_matches_triton now gives each impl its own buffer and asserts
both decode/prefill outputs are views into it. 42/42 pass.

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

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
(cherry picked from commit 06324dd3df)
2026-06-20 13:51:43 -07:00
Yongye ZhuandRoger Wang eb8e264edd [Model] M3 MSA indexer: unify top-k buffer, drop numpy, GPU-only decode plan
Follow-ups on the cudagraph-capturable MSA indexer:

- Top-k: both decode and prefill now write into the single shared, persistent
  topk_indices_buffer (decode at [:, :nd], prefill at [:, nd:]) and return views
  into it -- no fresh per-step top-k allocations.
- Build the decode plan + flat page table entirely with torch on-GPU: drop numpy
  and CpuGpuBuffer; segment offsets/lengths are computed via torch.cumsum into the
  persistent int32 buffers, and the request-major page table is scattered into the
  buffer via the on-GPU page indptr (the run bounds reads by indptr, so the full
  buffer is passed and no host page count is needed).
- No GPU->CPU sync on the decode path: scalars come from host ints
  (num_decode_tokens // num_decodes), and seq_lens.cpu() is confined to the eager
  prefill branch. The impl forward (fmha OnlyScore + Triton top-k) was already
  sync-free.

test_msa_indexer_impl_matches_triton now also asserts both outputs are views into
the persistent buffer. 42/42 in test_minimax_m3.py pass.

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

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
(cherry picked from commit d2fbaf73c1)
2026-06-20 13:51:43 -07:00
Yongye ZhuandRoger Wang 8f5070c447 [Model] Make MiniMax M3 MSA indexer cudagraph-capturable
The SM100 (MSA) lightning indexer declared AttentionCGSupport.NEVER and ran
eager: its fmha_sm100 score plan allocated fresh buffers every build() and the
run allocated a fresh max_score / top-k each call.

Make the decode side cudagraph-replay-safe:
- Reserve persistent plan buffers in the builder __init__ (sized over a scan of
  decode sizes [1, max_num_seqs]); build() fills them in place and calls the plan
  kernel directly with a fixed, batch-size-only num_kv_splits
  (estimate_num_kv_splits, a uniform-context replica of the planner's auto-split
  math; tunable via VLLM_M3_INDEXER_CONTEXT_LEN). A positive split count takes the
  deterministic plan path with no device->host sync.
- workspace_o / workspace_lse / cute_workspace are builder-owned dedicated tensors
  (not the shared global _alloc_workspace_buf cache) so a larger fmha call
  elsewhere can't realloc and move an address a captured graph baked.
- max_score is a 1-D-backed contiguous [H, max_k_tiles, nnz_qo] view (a sliced 3-D
  buffer is non-contiguous; the kernel assumes contiguous); max_k_tiles pinned for
  a stable shape.
- Top-k output goes to a model-level topk_indices_buffer (DeepSeek-V3.2 pattern),
  threaded model -> decoder layer -> sparse attention -> indexer; index_topk gains
  an out= param.
- Narrow _run_attention's eager break so the indexer runs in the captured segment
  and only the sparse attention is eager-broken. Builder reports UNIFORM_BATCH;
  prefill stays eager.

Verified: tests/kernels/attention/test_minimax_m3.py 42/42 pass (incl. impl-level
parity with num_kv_splits > 1 over short context); GPQA accuracy matches eager.

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

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

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
(cherry picked from commit 54810663ff)
2026-06-20 13:51:43 -07:00
Yongye ZhuandRoger Wang 934fa2b599 [Model] Add fmha_sm100 MSA indexer backend + fp8 index cache for MiniMax M3
Add an SM100/Blackwell lightning-indexer impl that computes the per-128-block
QK max-scores with fmha_sm100's score-only (OnlyScore) path and selects the
top-k blocks with the existing Triton minimax_m3_index_topk kernel, mirroring
how the main MSA attention pairs the SM100 attend with Triton. Decode and
prefill requests are split manually (decode-first batch) and each side gets its
own _fmha_sm100_plan / _fmha_sm100 call. Auto-selected on SM100 when
topk_blocks in (4, 8, 16, 32) for both bf16 and fp8 index caches; falls back to
the Triton indexer otherwise. The builder declares AttentionCGSupport.NEVER
(eager; the attention is broken out of the graph by _run_attention).

Extend the fused qknorm+rope+kv-insert kernel to optionally emit fp8 (e4m3) for
the index-K cache and index-Q via a direct cast with no scale tensors (RMSNorm
outputs are O(1) and scalar scales do not change top-k ordering). Only the index
outputs go fp8; q/k/v and q_out stay bf16 and bit-identical to the existing
path. MiniMaxM3IndexerCache now accepts fp8 caches and the model allocates
index_q in the cache dtype.

Tests: test_fmha_sm100_indexer_matches_reference (bf16/fp8 x prefill/decode) and
test_msa_indexer_impl_matches_triton (full impl parity vs the Triton indexer
through the real metadata builders); fp8 fused-kernel parity is covered in
test_fused_minimax_m3_qknorm_rope_kv_insert.

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

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

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
(cherry picked from commit bb4844dba3)
2026-06-20 13:51:43 -07:00
Jee Jee LiandRoger Wang 713fb6cdb7 Done (#87)
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
(cherry picked from commit 1d7ba0da2f0fedf056b0073fe7774f03fe4616f0)
2026-06-20 13:51:19 -07:00
Tyler Michael SmithGitHubClaudeCodexmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>kourosh hakhamaneshi
ebfbcfe46a Stop setting CUDA_VISIBLE_DEVICES internally in vLLM, add device_ids arg (#45026)
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Codex <codex@openai.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Co-authored-by: kourosh hakhamaneshi <kouroshHakha@users.noreply.github.com>
2026-06-20 13:38:10 -07:00
e9de72fe6c [Bugfix] Guard model_config access in _log_compilation_config (#46198)
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-06-20 19:26:38 +00:00
L丶GitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
d272418f45 [Perf] Optimize Qwen3-VL multi-video prompt processing (#46026)
Signed-off-by: Sirius29 <422058530@qq.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-06-20 07:09:18 -07:00
Sumanth R HegdeandGitHub 7ff7f5c8eb Revert "Fix Stale Encoder Cache After Weight Update" (#46125) 2026-06-20 07:09:09 -07:00
MattGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
dced290769 [Hardware][AMD][CI] Fix e2e core test group (#46024)
Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-06-20 02:04:35 -05:00
JasonLi314andGitHub 93bad11912 [Bugfix] Fix gridDim.y overflow for large row counts (#45255)
Signed-off-by: Jason Li <li.jason.cs@gmail.com>
2026-06-19 23:27:45 -04:00
djramicandGitHub 0fbf42af84 [ROCm] Fix VRAM not freed in test_phi3v (#46046)
Signed-off-by: Djordje Ramic <djoramic@amd.com>
2026-06-19 17:20:59 -05:00
Charlie FuandGitHub e6cd8913dd [ROCm][CI] Skip Qwen3.5-35B-A3B-MXFP4-AITER-TP2 for non gfx950 (#46109)
Signed-off-by: charlifu <charlifu@amd.com>
2026-06-19 17:20:10 -05:00
Ben BrowningandGitHub 859e4d436b [Bugfix][Parser] Fix U+FFFD leak at reasoning-to-content transition in engine parsers (#46159)
Signed-off-by: Ben Browning <bbrownin@redhat.com>
2026-06-19 22:09:28 +00:00
Micah WilliamsonandGitHub 4a083cc858 [ROCm][CI] Pin test_rocm_compressed_tensors_w8a8 to TRITON_ATTN (#46180)
Signed-off-by: Micah Williamson <micah.williamson@amd.com>
2026-06-19 15:20:06 -05:00
Vadim GimpelsonandGitHub ca7e1f2c43 Move CI failure diagnosis docs into ci-fails-buildkite skill (#45975)
Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com>
2026-06-19 20:12:40 +00:00
djramicandGitHub dec860fb19 [ROCm] Use vLLM's fp8 quant max in AITER hipBLASLt accuracy test (#46176)
Signed-off-by: Djordje Ramic <djoramic@amd.com>
2026-06-19 13:24:02 -05:00
Harry MellorandGitHub 0a49fb2b13 Fix dead link in docs (#46181)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
2026-06-19 18:16:09 +00:00
Ben BrowningandGitHub 4a8abf37c7 [Test] Migrate test_openai_schema.py to schemathesis 4.x (#46173)
Signed-off-by: Ben Browning <bbrownin@redhat.com>
2026-06-19 18:05:18 +00:00
01192139bf [DSv4] Pack KV caches into contiguous per-block allocations for DeepSeek V4 (#44577)
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>
Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>
Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-06-19 12:55:42 -04:00
Chris LeonardandGitHub b9a7cd464c [12/n] final _C library kernel migration (#45415) 2026-06-19 06:57:26 -07:00
69bdd34542 [Bugfix] Fall back to Pydantic loc for param in validation errors (#46038)
Signed-off-by: professorsab <135441198+professorsab@users.noreply.github.com>
Co-authored-by: Mahad Durrani <114791389+mahadrehmann@users.noreply.github.com>
2026-06-19 19:11:11 +08:00
Kunshang JiandGitHub ec67d7ae61 [xpu] bump up vllm-xpu-kernels v0.1.10 and upgrade 2618 umd (#40367)
Signed-off-by: Kunshang Ji <jikunshang95@gmail.com>
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
2026-06-19 15:37:20 +08:00
ecf9d83520 [AMD][CI] Fix Language Models Test (Extended Generation) failures (#45509)
Signed-off-by: Oxana Korzh <okorzh@amd.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-06-19 12:06:56 +08:00
Samuel ShenandGitHub c9135db27c [Docs] Update stale LMCache examples (#45762)
Signed-off-by: Samuel Shen <slshen@tensormesh.ai>
2026-06-19 03:21:36 +00:00
Jeff (Junze) MaGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2a6c6b9429 [DeepSeek-V4] Support TEP=16 for the block-FP8 shared expert (#46001)
Signed-off-by: Jeff Ma <jeffjma@umich.edu>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-06-18 20:10:12 -07:00
Jared WenandGitHub ab66606993 [bugfix]Indexer init skip and MTP TopK share for iteration (#45895)
Signed-off-by: JaredforReal <w13431838023@gmail.com>
2026-06-19 09:57:51 +08:00
9ea3a4015b [Bugfix] Fix corrupt outputs in MoE FP8 LoRA responses and MoE base model responses when LoRAs are loaded (#42120)
Signed-off-by: Nicholas Edelman <nedelman@nvidia.com>
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
Co-authored-by: Jee Jee Li <jeejeelee@inferact.ai>
Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>
2026-06-18 18:26:09 -07:00
Flora FengandGitHub 560fb8b867 [Cohere] Remove dead prepare_structured_tag override in Cohere parser (#46099)
Signed-off-by: sfeng33 <4florafeng@gmail.com>
2026-06-19 01:02:11 +00:00
Wentao YeandGitHub 675cd5d228 [Model Runner V2] Fix MRv2 memory leak test (#46095)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
2026-06-19 00:36:40 +00:00
7f616c327d [Bugfix] [Parser] Fix empty tool block silently dropping subsequent content (#46091)
Signed-off-by: Ben Browning <bbrownin@redhat.com>
Co-authored-by: Flora Feng <4florafeng@gmail.com>
2026-06-18 23:17:18 +00:00
Ivy XuandGitHub c3c6d723fd [Perf] Remove unused loggers in reasoning/ (#45988)
Signed-off-by: Ivy <fakeshadow1337@gmail.com>
2026-06-18 22:24:29 +00:00
Yifan QiaoGitHubJingyi Yangmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
41dcf49ca5 [Bugfix][KV Connector] Disable Mooncake TP put-striding when DCP > 1 (#45371)
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
Co-authored-by: Jingyi Yang <girasoleyang@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-06-18 15:13:44 -07:00
35e4dd4a69 [KV Connector][Mooncake] Async lookup to reduce scheduler overhead (#45659)
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>
2026-06-18 21:44:02 +00:00
4ce2d01453 fix(anthropic): auto-detect template support for mid-conversation system messages (#46025)
Signed-off-by: felix0080 <felix0080@users.noreply.github.com>
Signed-off-by: Ben Browning <bbrownin@redhat.com>
Co-authored-by: felix0080 <felix0080@users.noreply.github.com>
Co-authored-by: Ben Browning <bbrownin@redhat.com>
2026-06-18 16:19:11 -04:00
Woosuk KwonandGitHub 16908e132e [MRV2] Make FP32 Gumbel sampling more accurate (#45996)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
2026-06-18 19:42:09 +00:00
Wentao YeandGitHub 225936a1dd [CI Bug] Revert #42379 to fix CI Multi-Modal Models (Extended Generation 1) (#46070)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
2026-06-18 12:37:39 -07:00
f6ba720963 (security) Upgrade Starlette to >= 1.0.1 to fix CVE-2026-48710 (#45675)
Signed-off-by: jperezde <jperezde@redhat.com>
Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>
2026-06-18 12:35:13 -07:00
163 changed files with 4358 additions and 1114 deletions
+1 -1
View File
@@ -67,7 +67,7 @@ steps:
pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py &&
pytest -v -s v1/structured_output &&
pytest -v -s v1/test_serial_utils.py &&
pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py &&
pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py --ignore=v1/spec_decode/test_speculators_correctness.py &&
pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py --ignore=v1/kv_connector/unit/test_offloading_connector.py'
- label: "XPU server test"
depends_on:
+1 -14
View File
@@ -647,7 +647,7 @@ steps:
- pytest -v -s v1/cudagraph/test_cudagraph_mode.py
- label: e2e Core (1 GPU) # TBD
timeout_in_minutes: 180
timeout_in_minutes: 35
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
agent_pool: mi250_1
optional: true
@@ -2075,19 +2075,6 @@ steps:
- export VLLM_ALLOW_INSECURE_SERIALIZATION=1
- pytest -v -s v1/spec_decode/test_acceptance_length.py -m slow_test
- label: e2e Core (1 GPU) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
agent_pool: mi300_1
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/v1/
- tests/v1/e2e/
- vllm/platforms/rocm.py
commands:
- pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py
- label: e2e Scheduling (1 GPU) # TBD
timeout_in_minutes: 180
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+10
View File
@@ -74,6 +74,16 @@ steps:
- tests/v1/e2e/general/
commands:
- pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py
mirror:
amd:
device: mi250_1
timeout_in_minutes: 35
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/v1/
- tests/v1/e2e/general/
- vllm/platforms/rocm.py
- label: V1 e2e (2 GPUs)
key: v1-e2e-2-gpus
+1
View File
@@ -104,6 +104,7 @@ steps:
source_file_dependencies:
- csrc/quantization/
- vllm/model_executor/layers/quantization
- vllm/config/
- tests/kernels/quantization
- tests/kernels/quantization/test_rocm_skinny_gemms.py
- vllm/_aiter_ops.py
+9
View File
@@ -101,6 +101,15 @@ steps:
num_devices: 8
commands:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200.txt
mirror:
amd:
device: mi300_8
timeout_in_minutes: 180
depends_on:
- image-build-amd
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx.txt
- label: MoE Refactor Integration Test (H100 - TEMPORARY)
key: moe-refactor-integration-test-h100-temporary
@@ -68,7 +68,6 @@ steps:
- cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model # Otherwise, mp_method="spawn" doesn't work
mirror:
amd:
soft_fail: true
device: mi325_1
depends_on:
- image-build-amd
@@ -0,0 +1,35 @@
---
name: ci-fails-buildkite
description: Fetch and diagnose vLLM Buildkite CI failure logs. Use when investigating failing CI jobs on a PR or build, when the user pastes a buildkite.com URL, or asks to fetch/diagnose CI logs.
---
# Diagnosing vLLM Buildkite CI Failures
Buildkite logs are public; no login needed.
`.buildkite/scripts/ci-fetch-log.sh` saves each log as `ci-<build>-<job-name>.log`, stripped of timestamps and ANSI codes. Existing files are kept; set `CI_FETCH_LOG_FORCE=1` to refetch.
## Fetching logs
```bash
# All failed jobs in a PR's latest build (current branch's PR if omitted):
.buildkite/scripts/ci-fetch-log.sh --pr <PR>
# All failed jobs in a build (--soft also includes soft-failed jobs;
# --all fetches every finished job):
.buildkite/scripts/ci-fetch-log.sh "https://buildkite.com/vllm/ci/builds/<N>"
# One job — `gh pr checks` URLs (#<job_uuid>) and web UI URLs (?sid=) both
# work; pass "-" as a second argument to stream to stdout:
.buildkite/scripts/ci-fetch-log.sh "https://buildkite.com/vllm/ci/builds/<N>#<job_uuid>"
```
To clean an already-downloaded log with `.buildkite/scripts/ci-clean-log.sh`:
```bash
./ci-clean-log.sh ci.log
```
## Reference
See [docs/contributing/ci/failures.md](../../../docs/contributing/ci/failures.md) for the full guide: filing CI failure issues, investigating/bisecting, reproducing flaky tests, and daily triage.
+2 -3
View File
@@ -2,15 +2,14 @@
# for more info about CODEOWNERS file
# This lists cover the "core" components of vLLM that require careful review
/vllm/compilation @zou3519 @youkaichao @ProExpertProg @BoyuanFeng @vadiklyutiy
/vllm/compilation @zou3519 @youkaichao @ProExpertProg @BoyuanFeng
/vllm/distributed/kv_transfer @NickLucche @ApostaC @orozery @xuechendi
/vllm/lora @jeejeelee
/vllm/model_executor/layers/attention @LucasWilkinson @MatthewBonanni
/vllm/model_executor/layers/fused_moe @mgoin @pavanimajety @zyongye
/vllm/model_executor/layers/quantization @mgoin @robertgshaw2-redhat @tlrmchlsmth @yewentao256 @pavanimajety @zyongye
/vllm/model_executor/layers/mamba @tdoublep @tomeras91
/vllm/model_executor/layers/mamba/gdn_linear_attn.py @tdoublep @ZJY0516 @vadiklyutiy
/vllm/model_executor/layers/rotary_embedding.py @vadiklyutiy
/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @tdoublep @ZJY0516 @vadiklyutiy
/vllm/model_executor/model_loader @22quinn
/vllm/model_executor/layers/batch_invariant.py @yewentao256
/vllm/ir @ProExpertProg
+3 -1
View File
@@ -199,7 +199,9 @@ cython_debug/
.vscode/
# Claude
.claude/
.claude/*
!.claude/skills/
!.claude/skills/**
# Codex
.codex/
-11
View File
@@ -114,17 +114,6 @@ Follow these rules for all code changes in this repository:
- Keep comments and docstrings minimal and concise.
- Assume the reader is familiar with vLLM.
### Diagnosing CI failures
Buildkite logs are public; no login needed. Details: [docs/contributing/ci/failures.md](docs/contributing/ci/failures.md).
```bash
# All failed-job logs for a PR's latest build (current branch's PR if omitted):
.buildkite/scripts/ci-fetch-log.sh --pr <PR>
# Any Buildkite build or job URL also works:
.buildkite/scripts/ci-fetch-log.sh "<buildkite_url>"
```
### Commit messages
Add attribution using commit trailers such as `Co-authored-by:` (other projects use `Assisted-by:` or `Generated-by:`). For example:
+56 -72
View File
@@ -319,82 +319,35 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
endif()
#
# _C extension
# Legacy _C extension (ROCm only — CUDA ops migrated to _C_stable_libtorch)
#
set(VLLM_EXT_SRC
"csrc/quantization/activation_kernels.cu"
"csrc/torch_bindings.cpp")
if(VLLM_GPU_LANG STREQUAL "CUDA")
SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
# Set CUTLASS_REVISION. Used for FetchContent. Also fixes some bogus messages when building.
set(CUTLASS_REVISION "v4.4.2")
# Use the specified CUTLASS source directory for compilation if VLLM_CUTLASS_SRC_DIR is provided
if (DEFINED ENV{VLLM_CUTLASS_SRC_DIR})
set(VLLM_CUTLASS_SRC_DIR $ENV{VLLM_CUTLASS_SRC_DIR})
endif()
if(VLLM_CUTLASS_SRC_DIR)
if(NOT IS_ABSOLUTE VLLM_CUTLASS_SRC_DIR)
get_filename_component(VLLM_CUTLASS_SRC_DIR "${VLLM_CUTLASS_SRC_DIR}" ABSOLUTE)
endif()
message(STATUS "The VLLM_CUTLASS_SRC_DIR is set, using ${VLLM_CUTLASS_SRC_DIR} for compilation")
FetchContent_Declare(cutlass SOURCE_DIR ${VLLM_CUTLASS_SRC_DIR})
else()
FetchContent_Declare(
cutlass
GIT_REPOSITORY https://github.com/nvidia/cutlass.git
# Please keep this in sync with CUTLASS_REVISION line above.
GIT_TAG ${CUTLASS_REVISION}
GIT_PROGRESS TRUE
# Speed up CUTLASS download by retrieving only the specified GIT_TAG instead of the history.
# Important: If GIT_SHALLOW is enabled then GIT_TAG works only with branch names and tags.
# So if the GIT_TAG above is updated to a commit hash, GIT_SHALLOW must be set to FALSE
GIT_SHALLOW TRUE
)
endif()
FetchContent_MakeAvailable(cutlass)
set_gencode_flags_for_srcs(
SRCS "${VLLM_EXT_SRC}"
CUDA_ARCHS "${CUDA_ARCHS}")
# if CUDA endif
endif()
if (VLLM_GPU_LANG STREQUAL "HIP")
# Add QuickReduce kernels (ROCm-only; not part of stable ABI migration).
# TODO: Remove the cuda_view when ROCm upgrade to torch 2.11.
list(APPEND VLLM_EXT_SRC
if(VLLM_GPU_LANG STREQUAL "HIP")
set(VLLM_EXT_SRC
"csrc/torch_bindings.cpp"
"csrc/custom_quickreduce.cu"
"csrc/cuda_view.cu"
"csrc/libtorch_stable/cuda_utils_kernels.cu"
)
# if ROCM endif
endif()
"csrc/libtorch_stable/cuda_utils_kernels.cu")
message(STATUS "Enabling C extension.")
define_extension_target(
_C
DESTINATION vllm
LANGUAGE ${VLLM_GPU_LANG}
SOURCES ${VLLM_EXT_SRC}
COMPILE_FLAGS ${VLLM_GPU_FLAGS}
ARCHITECTURES ${VLLM_GPU_ARCHES}
INCLUDE_DIRECTORIES ${CUTLASS_INCLUDE_DIR}
INCLUDE_DIRECTORIES ${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}
USE_SABI 3
WITH_SOABI)
message(STATUS "Enabling C extension.")
define_extension_target(
_C
DESTINATION vllm
LANGUAGE ${VLLM_GPU_LANG}
SOURCES ${VLLM_EXT_SRC}
COMPILE_FLAGS ${VLLM_GPU_FLAGS}
ARCHITECTURES ${VLLM_GPU_ARCHES}
INCLUDE_DIRECTORIES ${CUTLASS_INCLUDE_DIR}
INCLUDE_DIRECTORIES ${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}
USE_SABI 3
WITH_SOABI)
# If CUTLASS is compiled on NVCC >= 12.5, it by default uses
# cudaGetDriverEntryPointByVersion as a wrapper to avoid directly calling the
# driver API. This causes problems when linking with earlier versions of CUDA.
# Setting this variable sidesteps the issue by calling the driver directly.
target_compile_definitions(_C PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
# If CUTLASS is compiled on NVCC >= 12.5, it by default uses
# cudaGetDriverEntryPointByVersion as a wrapper to avoid directly calling the
# driver API. This causes problems when linking with earlier versions of CUDA.
# Setting this variable sidesteps the issue by calling the driver directly.
target_compile_definitions(_C PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
endif() # _C HIP endif
if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
#
@@ -403,6 +356,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
set(VLLM_STABLE_EXT_SRC
"csrc/libtorch_stable/torch_bindings.cpp"
"csrc/libtorch_stable/activation_kernels.cu"
"csrc/libtorch_stable/quantization/activation_kernels.cu"
"csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu"
"csrc/libtorch_stable/quantization/w8a8/fp8/common.cu"
"csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu"
@@ -429,6 +383,38 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
"csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
if(VLLM_GPU_LANG STREQUAL "CUDA")
SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
# Set CUTLASS_REVISION. Used for FetchContent. Also fixes some bogus messages when building.
set(CUTLASS_REVISION "v4.4.2")
# Use the specified CUTLASS source directory for compilation if VLLM_CUTLASS_SRC_DIR is provided
if (DEFINED ENV{VLLM_CUTLASS_SRC_DIR})
set(VLLM_CUTLASS_SRC_DIR $ENV{VLLM_CUTLASS_SRC_DIR})
endif()
if(VLLM_CUTLASS_SRC_DIR)
if(NOT IS_ABSOLUTE VLLM_CUTLASS_SRC_DIR)
get_filename_component(VLLM_CUTLASS_SRC_DIR "${VLLM_CUTLASS_SRC_DIR}" ABSOLUTE)
endif()
message(STATUS "The VLLM_CUTLASS_SRC_DIR is set, using ${VLLM_CUTLASS_SRC_DIR} for compilation")
FetchContent_Declare(cutlass SOURCE_DIR ${VLLM_CUTLASS_SRC_DIR})
else()
FetchContent_Declare(
cutlass
GIT_REPOSITORY https://github.com/nvidia/cutlass.git
# Please keep this in sync with CUTLASS_REVISION line above.
GIT_TAG ${CUTLASS_REVISION}
GIT_PROGRESS TRUE
# Speed up CUTLASS download by retrieving only the specified GIT_TAG instead of the history.
# Important: If GIT_SHALLOW is enabled then GIT_TAG works only with branch names and tags.
# So if the GIT_TAG above is updated to a commit hash, GIT_SHALLOW must be set to FALSE
GIT_SHALLOW TRUE
)
endif()
FetchContent_MakeAvailable(cutlass)
list(APPEND VLLM_STABLE_EXT_SRC
"csrc/libtorch_stable/cuda_view.cu"
"csrc/libtorch_stable/cuda_utils_kernels.cu"
@@ -929,7 +915,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
SRCS "${FP4_SM120_SRCS}"
CUDA_ARCHS "${FP4_SM120_ARCHS}")
list(APPEND VLLM_STABLE_EXT_SRC "${FP4_SM120_SRCS}")
target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM120=1)
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1")
message(STATUS "Building SM12x NVFP4 for archs: ${FP4_SM120_ARCHS}")
@@ -962,7 +947,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
SRCS "${FP4_SM100_SRCS}"
CUDA_ARCHS "${FP4_SM100_ARCHS}")
list(APPEND VLLM_STABLE_EXT_SRC "${FP4_SM100_SRCS}")
target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM100=1)
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1")
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
message(STATUS "Building SM10x/11x NVFP4/MXFP4 for archs: ${FP4_SM100_ARCHS}")
+26 -1
View File
@@ -17,7 +17,7 @@ else()
FetchContent_Declare(
fmha_sm100
GIT_REPOSITORY https://github.com/vllm-project/MSA.git
GIT_TAG 544eee5e09ae2dfa774d5b06739013f9b7402c57
GIT_TAG fee783153f3efe57e3e933c5cb7e267a7cebcfb5
GIT_PROGRESS TRUE
CONFIGURE_COMMAND ""
BUILD_COMMAND ""
@@ -36,13 +36,38 @@ set(FMHA_SM100_PY_ROOT "${fmha_sm100_SOURCE_DIR}/python/fmha_sm100")
install(FILES
"${FMHA_SM100_PY_ROOT}/__init__.py"
"${FMHA_SM100_PY_ROOT}/api.py"
"${FMHA_SM100_PY_ROOT}/bench_utils.py"
"${FMHA_SM100_PY_ROOT}/jit.py"
"${FMHA_SM100_PY_ROOT}/sparse.py"
"${FMHA_SM100_PY_ROOT}/sparse_fmha_adapter.py"
DESTINATION vllm/third_party/fmha_sm100
COMPONENT fmha_sm100)
install(DIRECTORY "${FMHA_SM100_PY_ROOT}/csrc/"
DESTINATION vllm/third_party/fmha_sm100/csrc
COMPONENT fmha_sm100
PATTERN "__pycache__" EXCLUDE
PATTERN "*.pyc" EXCLUDE
PATTERN ".git*" EXCLUDE)
install(DIRECTORY "${FMHA_SM100_PY_ROOT}/cute/"
DESTINATION vllm/third_party/fmha_sm100/cute
COMPONENT fmha_sm100
PATTERN "__pycache__" EXCLUDE
PATTERN "*.pyc" EXCLUDE
PATTERN ".git*" EXCLUDE)
install(DIRECTORY "${FMHA_SM100_PY_ROOT}/cutlass/include/"
DESTINATION vllm/third_party/fmha_sm100/cutlass/include
COMPONENT fmha_sm100
PATTERN "__pycache__" EXCLUDE
PATTERN "*.pyc" EXCLUDE
PATTERN ".git*" EXCLUDE)
install(DIRECTORY "${FMHA_SM100_PY_ROOT}/cutlass/tools/util/include/"
DESTINATION vllm/third_party/fmha_sm100/cutlass/tools/util/include
COMPONENT fmha_sm100
PATTERN "__pycache__" EXCLUDE
PATTERN "*.pyc" EXCLUDE
PATTERN ".git*" EXCLUDE)
+29 -5
View File
@@ -60,6 +60,7 @@ endif()
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
set(QUTLASS_SOURCES
csrc/qutlass_registration.cpp
${qutlass_SOURCE_DIR}/qutlass/csrc/bindings.cpp
${qutlass_SOURCE_DIR}/qutlass/csrc/gemm.cu
${qutlass_SOURCE_DIR}/qutlass/csrc/gemm_ada.cu
@@ -78,8 +79,19 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
if(CUTLASS_INCLUDE_DIR AND EXISTS "${CUTLASS_INCLUDE_DIR}/cutlass/cutlass.h")
list(APPEND QUTLASS_INCLUDES "${CUTLASS_INCLUDE_DIR}")
if(CUTLASS_TOOLS_UTIL_INCLUDE_DIR AND
EXISTS "${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}/cutlass/util/packed_stride.hpp")
list(APPEND QUTLASS_INCLUDES "${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}")
else()
get_filename_component(_qutlass_cutlass_root "${CUTLASS_INCLUDE_DIR}" DIRECTORY)
if(EXISTS "${_qutlass_cutlass_root}/tools/util/include/cutlass/util/packed_stride.hpp")
list(APPEND QUTLASS_INCLUDES "${_qutlass_cutlass_root}/tools/util/include")
endif()
endif()
elseif(EXISTS "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include/cutlass/cutlass.h")
list(APPEND QUTLASS_INCLUDES "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include")
list(APPEND QUTLASS_INCLUDES
"${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include"
"${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/tools/util/include")
message(STATUS "[QUTLASS] Using QuTLASS vendored CUTLASS headers (no vLLM CUTLASS detected).")
else()
message(FATAL_ERROR "[QUTLASS] CUTLASS headers not found. "
@@ -91,12 +103,23 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
CUDA_ARCHS "${QUTLASS_ARCHS}"
)
target_sources(_C PRIVATE ${QUTLASS_SOURCES})
target_include_directories(_C PRIVATE ${QUTLASS_INCLUDES})
target_compile_definitions(_C PRIVATE
# QuTLASS uses legacy ATen headers and cannot be built with TORCH_TARGET_VERSION.
# Keep it as its own extension (registers torch.ops._qutlass_C).
define_extension_target(
_qutlass_C
DESTINATION vllm
LANGUAGE ${VLLM_GPU_LANG}
SOURCES ${QUTLASS_SOURCES}
COMPILE_FLAGS ${VLLM_GPU_FLAGS}
ARCHITECTURES ${VLLM_GPU_ARCHES}
INCLUDE_DIRECTORIES ${QUTLASS_INCLUDES}
USE_SABI 3
WITH_SOABI)
target_compile_definitions(_qutlass_C PRIVATE
QUTLASS_DISABLE_PYBIND=1
TARGET_CUDA_ARCH=${QUTLASS_TARGET_CC}
)
CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
set_property(SOURCE ${QUTLASS_SOURCES} APPEND PROPERTY COMPILE_OPTIONS
$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr --use_fast_math -O3>
@@ -111,4 +134,5 @@ else()
"[QUTLASS] Skipping build: no supported arch (12.0f / 10.0f) found in "
"CUDA_ARCHS='${CUDA_ARCHS}'.")
endif()
add_custom_target(_qutlass_C)
endif()
@@ -268,9 +268,14 @@ int64_t sm100_cutlass_mla_get_workspace_size(int64_t max_seq_len, int64_t num_ba
using TileShapeD = typename MlaSm100Type::TileShapeD;
arguments.problem_shape =
cute::make_tuple(TileShapeH{}, static_cast<int>(max_seq_len), TileShapeD{}, static_cast<int>(num_batches));
// Assumes device 0 when getting sm_count.
arguments.hw_info.sm_count =
sm_count <= 0 ? cutlass::KernelHardwareInfo::query_device_multiprocessor_count(/*device_id=*/0) : sm_count;
if (sm_count <= 0) {
int current_device = 0;
cudaGetDevice(&current_device);
arguments.hw_info.sm_count =
cutlass::KernelHardwareInfo::query_device_multiprocessor_count(current_device);
} else {
arguments.hw_info.sm_count = sm_count;
}
arguments.split_kv = static_cast<int>(num_kv_splits);
MlaSm100Type::Fmha::set_split_kv(arguments);
@@ -67,6 +67,13 @@
#include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh"
#endif
// Direct float -> E4M3 FP8 conversion for the indexer Q / index-K outputs.
#ifndef USE_ROCM
#include <cuda_fp8.h>
#else
#include <hip/hip_fp8.h>
#endif
#ifndef FINAL_MASK
#ifdef USE_ROCM
#define FINAL_MASK 0xffffffffffffffffULL
@@ -75,6 +82,19 @@
#endif
#endif
#ifdef USE_ROCM
// ROCm-compatible direct float -> E4M3 FP8 conversion (mirrors the DeepSeek V4
// fused kernel).
__device__ __forceinline__ uint8_t rocm_cvt_float_to_fp8_e4m3(float val) {
#if defined(HIP_FP8_TYPE_OCP)
__hip_fp8_e4m3 fp8_val(val);
#else
__hip_fp8_e4m3_fnuz fp8_val(val);
#endif
return reinterpret_cast<uint8_t&>(fp8_val);
}
#endif
namespace vllm {
namespace minimax_m3_fused_ops {
@@ -193,6 +213,8 @@ __device__ __forceinline__ void storeElems(
*reinterpret_cast<uint2*>(dst) = v;
}
// Main K/V cache store. kAuto = unquantized (cache_t == scalar_t); fp8 cache
// dtypes use the scaled-convert path with identity scale.
template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
__device__ __forceinline__ void storeCacheElems(
cache_t* __restrict__ dst, float const (&elems)[kElemsPerLane]) {
@@ -208,6 +230,32 @@ __device__ __forceinline__ void storeCacheElems(
}
}
// Store 4 fp32 registers -> 4 contiguous E4M3 FP8 bytes (direct cast,
// saturating to ±448). Used for the fp8 indexer-Q / index-K outputs; no scale
// (RMSNorm outputs are O(1) and the score path only needs relative block
// ordering).
__device__ __forceinline__ void storeElemsFp8(
uint8_t* __restrict__ dst, float const (&elems)[kElemsPerLane]) {
constexpr float kFp8Max = 448.0f;
#ifndef USE_ROCM
__nv_fp8x2_storage_t out2[kElemsPerLane / 2];
#pragma unroll
for (int i = 0; i < kElemsPerLane / 2; i++) {
float2 vv = make_float2(elems[2 * i], elems[2 * i + 1]);
vv.x = fminf(fmaxf(vv.x, -kFp8Max), kFp8Max);
vv.y = fminf(fmaxf(vv.y, -kFp8Max), kFp8Max);
out2[i] = __nv_cvt_float2_to_fp8x2(vv, __NV_SATFINITE, __NV_E4M3);
}
*reinterpret_cast<uint32_t*>(dst) = *reinterpret_cast<uint32_t const*>(out2);
#else
#pragma unroll
for (int i = 0; i < kElemsPerLane; i++) {
float vv = fminf(fmaxf(elems[i], -kFp8Max), kFp8Max);
dst[i] = rocm_cvt_float_to_fp8_e4m3(vv);
}
#endif
}
// ────────────────────────────────────────────────────────────────────────────
// Kernel
// ────────────────────────────────────────────────────────────────────────────
@@ -224,12 +272,14 @@ __device__ __forceinline__ void storeCacheElems(
// 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)
// cache_t/kv_dt: main attention KV-cache dtype (auto/fp8). out_idx_t/kFp8Idx:
// indexer index-K cache + index-Q output dtype (scalar_t or e4m3 byte).
template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt,
bool kIsSparse, bool kInsertKV>
typename out_idx_t, bool kIsSparse, bool kInsertKV, bool kFp8Idx>
__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* __restrict__ q_out, // [N, nq*128] contiguous, or nullptr
out_idx_t* __restrict__ index_q_out, // [N, niq*128]; scalar_t or e4m3 byte
scalar_t const* __restrict__ q_norm_w,
scalar_t const* __restrict__ k_norm_w,
scalar_t const* __restrict__ iq_norm_w,
@@ -238,8 +288,8 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
int64_t const* __restrict__ positions, // [N] i64
int64_t const* __restrict__ slot_mapping, // main K/V slots or nullptr
int64_t const* __restrict__ index_slot_mapping, // index K slots/nullptr
cache_t* __restrict__ kv_cache, // [nb,2,bs,nkv,128] or nullptr
scalar_t* __restrict__ index_cache, // [nb*bs, 128] or nullptr
cache_t* __restrict__ kv_cache, // [nb,2,bs,nkv,128] or nullptr
out_idx_t* __restrict__ index_cache, // [nb*bs, 128]; scalar_t or e4m3 byte
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].
@@ -334,9 +384,12 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
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;
// bf16 index_q_out: gather here. fp8: written by the explicit fp8 store.
if constexpr (!kFp8Idx) {
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
@@ -356,7 +409,19 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
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);
if constexpr (kFp8Idx) {
// index_q is e4m3 bytes; Q/K (and in-place index_k) stay scalar_t.
if (isIQ && index_q_out != nullptr) {
storeElemsFp8(index_q_out +
static_cast<int64_t>(tokenIdx) * niq * kHeadDim +
(slot - iq_begin) * kHeadDim + dim_base,
elems);
} else {
storeElems<scalar_t>(store_ptr + dim_base, elems);
}
} else {
storeElems<scalar_t>(store_ptr + dim_base, elems);
}
}
// ── Cache inserts (sparse serving only). ───────────────────────────────
@@ -367,8 +432,11 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
: (isIK ? index_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);
if constexpr (kFp8Idx) {
storeElemsFp8(index_cache + sm * kHeadDim + dim_base, elems);
} else {
storeElems<scalar_t>(index_cache + sm * kHeadDim + dim_base, 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);
@@ -398,19 +466,19 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
// Launch wrapper
// ────────────────────────────────────────────────────────────────────────────
template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
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,
int64_t const* index_slot_mapping, cache_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) {
void launchFusedMiniMaxM3(
scalar_t* qkv, scalar_t* q_out, void* 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,
int64_t const* index_slot_mapping, cache_t* kv_cache, void* 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,
bool const fp8_idx, cudaStream_t stream) {
// Index outputs are scalar_t (bf16) or e4m3 bytes (uint8_t); reinterpret the
// void* pointers per instantiation in the LAUNCH macro.
// 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;
@@ -440,25 +508,27 @@ void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
config.attrs = attrs;
config.numAttrs = (sm_version >= 90) ? 1 : 0;
#define LAUNCH(IS_SPARSE, INSERT) \
cudaLaunchKernelEx( \
&config, \
fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt, \
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, index_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)
#define LAUNCH(IS_SPARSE, INSERT, FP8, OUT_T) \
cudaLaunchKernelEx( \
&config, \
fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt, OUT_T, \
IS_SPARSE, INSERT, FP8>, \
qkv, q_out, reinterpret_cast<OUT_T*>(index_q_out), q_norm_w, k_norm_w, \
iq_norm_w, ik_norm_w, cos_sin_cache, positions, slot_mapping, \
index_slot_mapping, kv_cache, reinterpret_cast<OUT_T*>(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, cache_t, kv_dt, \
IS_SPARSE, INSERT> \
#define LAUNCH(IS_SPARSE, INSERT, FP8, OUT_T) \
fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt, OUT_T, \
IS_SPARSE, INSERT, FP8> \
<<<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, \
index_slot_mapping, kv_cache, index_cache, eps, rotary_dim, \
qkv, q_out, reinterpret_cast<OUT_T*>(index_q_out), q_norm_w, \
k_norm_w, iq_norm_w, ik_norm_w, cos_sin_cache, positions, \
slot_mapping, index_slot_mapping, kv_cache, \
reinterpret_cast<OUT_T*>(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
@@ -466,14 +536,22 @@ void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
if (has_index) {
if (insert_kv) {
LAUNCH(true, true); // sparse serving
if (fp8_idx) {
LAUNCH(true, true, true, uint8_t); // sparse serving, fp8 index outputs
} else {
LAUNCH(true, true, false, scalar_t); // sparse serving, bf16
}
} else {
LAUNCH(true, false); // sparse profiling
if (fp8_idx) {
LAUNCH(true, false, true, uint8_t); // sparse profiling, fp8 index_q
} else {
LAUNCH(true, false, false, scalar_t); // sparse profiling, bf16
}
}
} else {
// Dense layer: never has an index branch and never inserts here (the
// generic Attention layer owns the KV insert).
LAUNCH(false, false);
LAUNCH(false, false, false, scalar_t);
}
#undef LAUNCH
}
@@ -485,8 +563,9 @@ void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
vllm::minimax_m3_fused_ops::launchFusedMiniMaxM3<st, CACHE_T, KV_DTYPE>( \
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, \
index_q_out.has_value() \
? reinterpret_cast<void*>(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()) \
@@ -502,11 +581,11 @@ void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
: nullptr, \
insert_kv ? reinterpret_cast<CACHE_T*>(kv_cache->data_ptr()) : nullptr, \
(insert_kv && has_index) \
? reinterpret_cast<st*>(index_cache->data_ptr()) \
? reinterpret_cast<void*>(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)
kv_s_head, has_index, insert_kv, fp8_idx, stream)
// ────────────────────────────────────────────────────────────────────────────
// Torch op wrapper
@@ -612,6 +691,7 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
torch::headeronly::ScalarType::Long &&
index_slot_mapping->numel() == slot_mapping->numel()),
"index_slot_mapping must be int64 CUDA with slot_mapping length");
// Main attention KV cache: auto matches qkv, fp8 uses uint8 storage.
if (kv_dt == vllm::Fp8KVCacheDataType::kAuto) {
STD_TORCH_CHECK(kv_cache->scalar_type() == qkv.scalar_type(),
"auto kv_cache dtype must match qkv");
@@ -620,9 +700,13 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
kv_cache->scalar_type() == torch::headeronly::ScalarType::Byte,
"fp8 kv_cache must use uint8 storage");
}
STD_TORCH_CHECK(index_cache.has_value() &&
index_cache->scalar_type() == qkv.scalar_type(),
"insert mode requires matching index_cache");
// Indexer index-K cache: independent dtype -- qkv dtype or fp8 e4m3.
STD_TORCH_CHECK(
index_cache.has_value() &&
(index_cache->scalar_type() == qkv.scalar_type() ||
index_cache->scalar_type() ==
torch::headeronly::ScalarType::Float8_e4m3fn),
"insert mode requires index_cache matching qkv dtype or fp8 e4m3");
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)");
@@ -652,14 +736,31 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
"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");
(index_q_out->scalar_type() == qkv.scalar_type() ||
index_q_out->scalar_type() ==
torch::headeronly::ScalarType::Float8_e4m3fn),
"index_q_out must be contiguous CUDA, qkv dtype or fp8 e4m3");
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");
}
// fp8 index path: the index-K cache and index-Q outputs are e4m3 bytes while
// q/k/v + q_out stay qkv dtype. Both index outputs must agree.
auto const kFp8 = torch::headeronly::ScalarType::Float8_e4m3fn;
bool const fp8_idx =
(index_cache.has_value() && index_cache->scalar_type() == kFp8) ||
(index_q_out.has_value() && index_q_out->scalar_type() == kFp8);
if (fp8_idx) {
STD_TORCH_CHECK(
!index_cache.has_value() || index_cache->scalar_type() == kFp8,
"fp8 index path: index_cache must be fp8 e4m3");
STD_TORCH_CHECK(
!index_q_out.has_value() || index_q_out->scalar_type() == kFp8,
"fp8 index path: index_q_out must be fp8 e4m3");
}
const torch::stable::accelerator::DeviceGuard device_guard(
qkv.get_device_index());
auto stream = get_current_cuda_stream(qkv.get_device_index());
+7 -6
View File
@@ -81,11 +81,11 @@ __global__ void rms_norm_kernel(
#pragma unroll
for (int j = 0; j < VEC_SIZE; j++) {
float x = static_cast<float>(src1.val[j]);
scalar_t normalized = static_cast<scalar_t>(x * s_variance);
if constexpr (HasWeight) {
dst.val[j] = normalized * src2.val[j];
float w = static_cast<float>(src2.val[j]);
dst.val[j] = static_cast<scalar_t>(x * s_variance * w);
} else {
dst.val[j] = normalized;
dst.val[j] = static_cast<scalar_t>(x * s_variance);
}
}
v_out[i] = dst;
@@ -151,7 +151,8 @@ fused_add_rms_norm_kernel(
#pragma unroll
for (int j = 0; j < width; ++j) {
float x = Converter::convert(res.data[j]);
out.data[j] = Converter::convert(x * s_variance) * w.data[j];
float wf = Converter::convert(w.data[j]);
out.data[j] = Converter::convert(x * s_variance * wf);
}
} else {
#pragma unroll
@@ -198,8 +199,8 @@ fused_add_rms_norm_kernel(
for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
float x = (float)residual[blockIdx.x * hidden_size + idx];
if constexpr (HasWeight) {
input[blockIdx.x * input_stride + idx] =
(scalar_t)(x * s_variance) * weight[idx];
float w = (float)weight[idx];
input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance * w);
} else {
input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance);
}
@@ -66,8 +66,13 @@ __global__ void rms_norm_static_fp8_quant_kernel(
#pragma unroll
for (int j = 0; j < VEC_SIZE; j++) {
float x = static_cast<float>(src1.val[j]);
// Multiply in weight's native dtype to match rms_norm_kernel.
scalar_t out_norm = static_cast<scalar_t>(x * s_variance) * src2.val[j];
float w = static_cast<float>(src2.val[j]);
// Round normalized result through scalar_t to match the precision of the
// unfused composite (rms_norm writes scalar_t, then
// static_scaled_fp8_quant re-loads it as float before FP8 conversion).
// Without this round, the fused path is strictly more accurate and
// disagrees with the composite at exact E4M3 quantization tie boundaries.
scalar_t out_norm = static_cast<scalar_t>(x * s_variance * w);
out[blockIdx.x * hidden_size + idx * VEC_SIZE + j] =
scaled_fp8_conversion<true, fp8_type>(static_cast<float>(out_norm),
scale_inv);
@@ -137,8 +142,12 @@ fused_add_rms_norm_static_fp8_quant_kernel(
#pragma unroll
for (int i = 0; i < width; ++i) {
float x = Converter::convert(res.data[i]);
// Multiply in weight's native dtype to match fused_add_rms_norm_kernel.
HipT out_norm_h = Converter::convert(x * s_variance) * w.data[i];
float wf = Converter::convert(w.data[i]);
// See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
// to match the unfused composite path at FP8 boundaries. We use the
// backend's hip_type for the intermediate since c10::Half/BFloat16 has
// ambiguous conversions on CUDA and no implicit conversion on ROCm.
HipT out_norm_h = Converter::convert(x * s_variance * wf);
out[id * width + i] = scaled_fp8_conversion<true, fp8_type>(
Converter::convert(out_norm_h), scale_inv);
}
@@ -183,8 +192,10 @@ fused_add_rms_norm_static_fp8_quant_kernel(
for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
float x = (float)residual[blockIdx.x * hidden_size + idx];
// Multiply in weight's native dtype to match fused_add_rms_norm_kernel.
scalar_t out_norm = static_cast<scalar_t>(x * s_variance) * weight[idx];
float w = (float)weight[idx];
// See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
// to match the unfused composite path at FP8 boundaries.
scalar_t out_norm = static_cast<scalar_t>(x * s_variance * w);
out[blockIdx.x * hidden_size + idx] = scaled_fp8_conversion<true, fp8_type>(
static_cast<float>(out_norm), scale_inv);
}
@@ -9,7 +9,7 @@
#include <torch/headeronly/core/ScalarType.h>
#include "../../cuda_compat.h"
#include "core/math.hpp"
#include "libtorch_stable/core/math.hpp"
#include "libtorch_stable/dispatch_utils.h"
#include "libtorch_stable/torch_utils.h"
+28
View File
@@ -2,9 +2,25 @@
#include <torch/csrc/stable/library.h>
#include <torch/csrc/stable/tensor.h>
#include <torch/headeronly/util/Exception.h>
#include <optional>
#include <string>
#include <vector>
#include <torch/csrc/stable/ops.h>
inline torch::stable::Tensor weak_ref_tensor(torch::stable::Tensor& tensor) {
// Ensure tensor is on CUDA
STD_TORCH_CHECK(tensor.device().is_cuda(), "Tensor must be on CUDA device");
// Get the raw data pointer
void* data_ptr = tensor.mutable_data_ptr();
/// Create a new tensor from the raw data pointer
return torch::stable::from_blob(data_ptr, tensor.sizes(), tensor.strides(),
tensor.device(), tensor.scalar_type());
}
void per_token_group_quant_fp8(const torch::stable::Tensor& input,
torch::stable::Tensor& output_q,
@@ -371,6 +387,18 @@ 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,
double alpha = 1.0, double beta = 0.0);
void silu_and_mul_quant(torch::stable::Tensor& out,
torch::stable::Tensor& input,
torch::stable::Tensor& scale);
void persistent_masked_m_silu_mul_quant(
const torch::stable::Tensor& input, // (E, T, 2*H)
const torch::stable::Tensor& tokens_per_expert, // (E)
torch::stable::Tensor& y_q, // (E, T, H) [OUT]
torch::stable::Tensor& y_s, // (E, T, H//group_size) [OUT]
bool use_ue8m0);
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,
@@ -1,16 +1,12 @@
#include <ATen/cuda/CUDAContext.h>
#include <torch/all.h>
#include <c10/cuda/CUDAGuard.h>
#include "libtorch_stable/torch_utils.h"
#include <cmath>
#include "core/math.hpp"
#include "../cuda_compat.h"
#include "dispatch_utils.h"
#include "libtorch_stable/core/math.hpp"
#include "cuda_compat.h"
#include "libtorch_stable/dispatch_utils.h"
#include "quantization/w8a8/fp8/common.cuh"
#include <c10/util/Float8_e4m3fn.h>
#ifndef USE_ROCM
#include <cuda_bf16.h>
#include <cuda_fp16.h>
@@ -33,7 +29,6 @@ typedef __hip_fp8x4_e4m3_fnuz __nv_fp8x4_e4m3;
#endif
#endif
#include "core/registration.h"
namespace vllm {
template <typename T>
@@ -564,41 +559,47 @@ __global__ void silu_mul_fp8_quant_deep_gemm_kernel(
} // namespace vllm
// Launch activation, gating, and quantize kernel.
#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL) \
int d = input.size(-1) / 2; \
int64_t num_tokens = input.numel() / input.size(-1); \
dim3 grid(num_tokens, num_tokens > 16 ? num_tokens > 32 ? 1 : 2 : 4); \
dim3 block(std::min(d, 512)); \
const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \
const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); \
VLLM_DISPATCH_FLOATING_TYPES( \
input.scalar_type(), "act_and_mul_kernel", [&] { \
VLLM_DISPATCH_FP8_TYPES( \
out.scalar_type(), "fused_add_rms_norm_kernel_fp8_type", [&] { \
vllm::act_and_mul_quant_kernel<scalar_t, KERNEL<scalar_t>, \
fp8_t> \
<<<grid, block, 0, stream>>>(out.data_ptr<fp8_t>(), \
input.data_ptr<scalar_t>(), \
scale.data_ptr<float>(), d); \
}); \
#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL) \
int d = input.size(-1) / 2; \
int64_t num_tokens = input.numel() / input.size(-1); \
dim3 grid(num_tokens, num_tokens > 16 ? num_tokens > 32 ? 1 : 2 : 4); \
dim3 block(std::min(d, 512)); \
const torch::stable::accelerator::DeviceGuard device_guard( \
input.get_device_index()); \
const cudaStream_t stream = \
get_current_cuda_stream(input.get_device_index()); \
VLLM_STABLE_DISPATCH_FLOATING_TYPES( \
input.scalar_type(), "act_and_mul_kernel", [&] { \
VLLM_STABLE_DISPATCH_FP8_TYPES( \
out.scalar_type(), "act_and_mul_quant_kernel_fp8_type", [&] { \
vllm::act_and_mul_quant_kernel<scalar_t, KERNEL<scalar_t>, \
fp8_t> \
<<<grid, block, 0, stream>>>( \
out.mutable_data_ptr<fp8_t>(), \
input.const_data_ptr<scalar_t>(), \
scale.const_data_ptr<float>(), d); \
}); \
});
void silu_and_mul_quant(torch::Tensor& out, // [..., d]
torch::Tensor& input, // [..., 2 * d]
torch::Tensor& scale) {
TORCH_CHECK(out.dtype() == torch::kFloat8_e4m3fn ||
out.dtype() == torch::kFloat8_e4m3fnuz);
TORCH_CHECK(input.dtype() == torch::kFloat16 ||
input.dtype() == torch::kBFloat16);
TORCH_CHECK(input.size(-1) % 2 == 0);
void silu_and_mul_quant(torch::stable::Tensor& out, // [..., d]
torch::stable::Tensor& input, // [..., 2 * d]
torch::stable::Tensor& scale) {
STD_TORCH_CHECK(
out.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn ||
out.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fnuz);
STD_TORCH_CHECK(
input.scalar_type() == torch::headeronly::ScalarType::Half ||
input.scalar_type() == torch::headeronly::ScalarType::BFloat16,
"Input must be FP16 or BF16");
STD_TORCH_CHECK(input.size(-1) % 2 == 0);
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel);
}
void persistent_masked_m_silu_mul_quant(
const at::Tensor& input, // (E, T, 2*H)
const at::Tensor& tokens_per_expert, // (E)
at::Tensor& y_q, // (E, T, H) [OUT]
at::Tensor& y_s, // (E, T, H//group_size) [OUT]
const torch::stable::Tensor& input, // (E, T, 2*H)
const torch::stable::Tensor& tokens_per_expert, // (E)
torch::stable::Tensor& y_q, // (E, T, H) [OUT]
torch::stable::Tensor& y_s, // (E, T, H//group_size) [OUT]
bool cast_scale_ue8m0) {
#ifndef USE_ROCM
@@ -606,14 +607,18 @@ void persistent_masked_m_silu_mul_quant(
// fixed GROUP_SIZE of 128.
static constexpr int GROUP_SIZE = 128;
TORCH_CHECK(input.dtype() == torch::kBFloat16);
TORCH_CHECK(y_q.dtype() == torch::kFloat8_e4m3fn ||
y_q.dtype() == torch::kFloat8_e4m3fnuz);
TORCH_CHECK(input.size(-1) % (GROUP_SIZE * 2) == 0);
STD_TORCH_CHECK(input.scalar_type() ==
torch::headeronly::ScalarType::BFloat16);
STD_TORCH_CHECK(
y_q.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn ||
y_q.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fnuz);
STD_TORCH_CHECK(input.size(-1) % (GROUP_SIZE * 2) == 0);
bool const is_packed_ue8m0 =
(y_s.dtype() == torch::kInt32 && cast_scale_ue8m0);
TORCH_CHECK(y_s.dtype() == torch::kFloat32 || is_packed_ue8m0);
(y_s.scalar_type() == torch::headeronly::ScalarType::Int &&
cast_scale_ue8m0);
STD_TORCH_CHECK(y_s.scalar_type() == torch::headeronly::ScalarType::Float ||
is_packed_ue8m0);
using Idx_t = int64_t;
@@ -631,7 +636,7 @@ void persistent_masked_m_silu_mul_quant(
int const NUM_GROUPS = H / GROUP_SIZE;
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
// TODO: Get this from cuda_arch ?
static constexpr int SILU_V2_BLOCK_COUNT = 132 * 32;
@@ -643,18 +648,21 @@ void persistent_masked_m_silu_mul_quant(
static constexpr int max_shared_mem_bytes = \
GROUP_SIZE * 2 * STAGES * NUM_WARPS * 2; \
dim3 grid(sms), block(THREAD_COUNT); \
const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \
VLLM_DISPATCH_FP8_TYPES( \
const torch::stable::accelerator::DeviceGuard device_guard( \
input.get_device_index()); \
VLLM_STABLE_DISPATCH_FP8_TYPES( \
y_q.scalar_type(), "silu_mul_fp8_quant_deep_gemm_kernel", [&] { \
vllm::silu_mul_fp8_quant_deep_gemm_kernel< \
BLOCK_COUNT, max_shared_mem_bytes, fp8_t, scale_t, THREAD_COUNT, \
Idx_t, CEIL_UE8M0, GROUP_SIZE, STAGES> \
<<<grid, block, max_shared_mem_bytes + (E + 1) * 16, stream>>>( \
reinterpret_cast<__nv_bfloat16*>(input.data_ptr()), \
(fp8_t*)y_q.data_ptr(), \
reinterpret_cast<scale_t*>(y_s.data_ptr()), \
reinterpret_cast<int32_t*>(tokens_per_expert.data_ptr()), E, \
T, H, stride_i_e, stride_i_t, stride_i_h, stride_yq_e, \
reinterpret_cast<const __nv_bfloat16*>( \
input.const_data_ptr()), \
y_q.mutable_data_ptr<fp8_t>(), \
reinterpret_cast<scale_t*>(y_s.mutable_data_ptr()), \
reinterpret_cast<const int32_t*>( \
tokens_per_expert.const_data_ptr()), \
E, T, H, stride_i_e, stride_i_t, stride_i_h, stride_yq_e, \
stride_yq_t, stride_yq_h, STRIDE_YS_E, STRIDE_YS_T, \
STRIDE_YS_G, STRIDE_YS_P, stride_counts_e); \
});
@@ -679,7 +687,7 @@ void persistent_masked_m_silu_mul_quant(
Idx_t stride_ys_g = y_s.stride(2);
Idx_t stride_ys_p = 0;
if (!cast_scale_ue8m0) {
TORCH_CHECK(!is_packed_ue8m0);
STD_TORCH_CHECK(!is_packed_ue8m0);
LAUNCH_ON_H(float, stride_ys_e, stride_ys_t, stride_ys_g, stride_ys_p,
false);
return;
@@ -692,8 +700,8 @@ void persistent_masked_m_silu_mul_quant(
return;
}
TORCH_CHECK(cast_scale_ue8m0 && is_packed_ue8m0);
TORCH_CHECK(y_s.dtype() == torch::kInt32);
STD_TORCH_CHECK(cast_scale_ue8m0 && is_packed_ue8m0);
STD_TORCH_CHECK(y_s.scalar_type() == torch::headeronly::ScalarType::Int);
// Int32 packed ue8m0 scales tensor.
// Let E, T, G be the number to experts, number of tokens and number of groups
@@ -31,7 +31,7 @@
#include "cutlass/util/packed_stride.hpp"
#include "core/math.hpp"
#include "libtorch_stable/core/math.hpp"
#include "core/batch_invariant.hpp"
using namespace cute;
@@ -31,7 +31,7 @@
#include "cutlass/util/packed_stride.hpp"
#include "core/math.hpp"
#include "libtorch_stable/core/math.hpp"
#include "core/batch_invariant.hpp"
using namespace cute;
@@ -19,7 +19,7 @@
#include "cutlass/gemm/collective/collective_builder.hpp"
#include "cutlass/util/packed_stride.hpp"
#include "core/math.hpp"
#include "libtorch_stable/core/math.hpp"
#include "libtorch_stable/cutlass_extensions/common.hpp"
// clang-format on
@@ -14,7 +14,7 @@
#include "cutlass/epilogue/collective/collective_builder.hpp"
#include "cutlass/gemm/collective/collective_builder.hpp"
#include "core/math.hpp"
#include "libtorch_stable/core/math.hpp"
#include "libtorch_stable/cutlass_extensions/common.hpp"
// clang-format on
@@ -22,7 +22,7 @@
#include "cutlass/epilogue/threadblock/fusion/visitors.hpp"
#include "cutlass/gemm/kernel/default_gemm_universal_with_visitor.h"
#include "core/math.hpp"
#include "libtorch_stable/core/math.hpp"
#include "libtorch_stable/cutlass_extensions/common.hpp"
// clang-format on
@@ -301,8 +301,9 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel(
const int sf_k_local = local_group_id % kGroupsPerBlockX;
const int row_local = local_group_id / kGroupsPerBlockX;
const int sf_k_idx = blockIdx.x * kGroupsPerBlockX + sf_k_local;
const int mn_idx = blockIdx.y * kRowsPerBlock + row_local;
// Rows on grid.x: mn scales with tokens and can exceed the 65535 grid.y cap.
const int sf_k_idx = blockIdx.y * kGroupsPerBlockX + sf_k_local;
const int mn_idx = blockIdx.x * kRowsPerBlock + row_local;
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
asm volatile("griddepcontrol.wait;");
@@ -496,14 +497,15 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
" is not a multiple of 4.");
const int kx = GetGroupsPerBlockX(padded_groups_per_row);
const int ry = 16 / kx;
const int64_t blocks_x = padded_groups_per_row / kx;
const int64_t blocks_y = (tma_aligned_mn + ry - 1) / ry;
const int64_t row_blocks = (tma_aligned_mn + ry - 1) / ry;
const int64_t sf_k_blocks = padded_groups_per_row / kx;
const int num_threads = (kx * ry) * THREADS_PER_GROUP;
// CUDA caps grid.x and grid.y at 2^31 - 1; guard against pathological inputs.
STD_TORCH_CHECK(blocks_x <= static_cast<int64_t>(INT32_MAX) &&
blocks_y <= static_cast<int64_t>(INT32_MAX),
// CUDA caps grid.x at 2^31 - 1 and grid.y at 2^16 - 1 (65535).
constexpr int64_t kMaxGridDimYZ = 65535;
STD_TORCH_CHECK(row_blocks <= static_cast<int64_t>(INT32_MAX) &&
sf_k_blocks <= kMaxGridDimYZ,
"per_token_group_quant_8bit_packed grid too large: (",
blocks_x, ", ", blocks_y, ").");
row_blocks, ", ", sf_k_blocks, ").");
auto dst_type = output_q.scalar_type();
@@ -513,8 +515,8 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
#define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, KX, RY) \
do { \
cudaLaunchConfig_t config = {}; \
config.gridDim = dim3(static_cast<unsigned int>(blocks_x), \
static_cast<unsigned int>(blocks_y)); \
config.gridDim = dim3(static_cast<unsigned int>(row_blocks), \
static_cast<unsigned int>(sf_k_blocks)); \
config.blockDim = dim3(num_threads); \
config.dynamicSmemBytes = 0; \
config.stream = stream; \
@@ -539,8 +541,8 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
#else
#define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, KX, RY) \
do { \
dim3 grid(static_cast<unsigned int>(blocks_x), \
static_cast<unsigned int>(blocks_y)); \
dim3 grid(static_cast<unsigned int>(row_blocks), \
static_cast<unsigned int>(sf_k_blocks)); \
dim3 block(num_threads); \
per_token_group_quant_8bit_packed_register_kernel<T, DST_DTYPE, 128, KX, \
RY> \
+27
View File
@@ -34,6 +34,20 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
// TODO: Remove this once ROCm upgrade to torch 2.11.
ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
// Note about marlin kernel 'workspace' arguments:
// Technically these should be mutable since they are modified by the kernel.
// But since they are set back to zero once the kernel is finished we can
// hand wave and say that they have no net effect.
//
// The reason to mark 'workspace' as immutable is so that they don't interfere
// with using ScalarType arguments in the ops. If they are marked as mutable,
// pytorch throws an assert in
// 'torch._higher_order_ops._register_effectful_op' that prevents these
// kernels from being torch.compile'd.
// See the following document for more info on custom types and ops that use
// custom types:
// https://docs.google.com/document/d/18fBMPuOJ0fY5ZQ6YyrHUppw9FA332CpNtgB6SOIgyuA
// Machete (Dense) Optimized Mixed Precision GEMM for Hopper.
ops.def(
"machete_supported_schedules("
@@ -480,6 +494,11 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
"Tensor workspace, int k, int max_seq_len) -> ()");
// Activation ops
ops.def(
"persistent_masked_m_silu_mul_quant(Tensor input, Tensor counts, Tensor! "
"y_q, Tensor! y_s, bool use_ue8m0) -> ()");
ops.def("weak_ref_tensor(Tensor input) -> Tensor");
// Activation function used in SwiGLU.
ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()");
@@ -492,6 +511,10 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
"silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit, "
"float alpha=1.0, float beta=0.0) -> ()");
// SwiGLU activation with FP8 quantization.
ops.def(
"silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
// Activation function used in GeGLU with `none` approximation.
ops.def("gelu_and_mul(Tensor! out, Tensor input) -> ()");
@@ -690,6 +713,10 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
ops.impl("persistent_topk", TORCH_BOX(&persistent_topk));
// Activation kernels (shared CUDA/ROCm)
ops.impl("persistent_masked_m_silu_mul_quant",
TORCH_BOX(&persistent_masked_m_silu_mul_quant));
ops.impl("weak_ref_tensor", TORCH_BOX(&weak_ref_tensor));
ops.impl("silu_and_mul_quant", TORCH_BOX(&silu_and_mul_quant));
ops.impl("silu_and_mul", TORCH_BOX(&silu_and_mul));
ops.impl("mul_and_silu", TORCH_BOX(&mul_and_silu));
ops.impl("gelu_and_mul", TORCH_BOX(&gelu_and_mul));
-32
View File
@@ -9,28 +9,6 @@
#include <vector>
torch::Tensor weak_ref_tensor(torch::Tensor& tensor) {
// Ensure tensor is on CUDA
if (!tensor.is_cuda()) {
throw std::runtime_error("Tensor must be on CUDA device");
}
// Get the raw data pointer
void* data_ptr = tensor.data_ptr();
// Get tensor sizes and strides
std::vector<int64_t> sizes = tensor.sizes().vec();
std::vector<int64_t> strides = tensor.strides().vec();
// Get tensor options (dtype, device)
auto options = tensor.options();
// Create a new tensor from the raw data pointer
auto new_tensor = torch::from_blob(data_ptr, sizes, strides, options);
return new_tensor;
}
// rms_norm and fused_add_rms_norm declarations also exist in
// csrc/libtorch_stable/ops.h (torch::stable ABI for CUDA). They remain here
// because the CPU build still uses these torch::Tensor declarations.
@@ -53,16 +31,6 @@ void silu_and_mul(torch::Tensor& out, torch::Tensor& input);
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);
void persistent_masked_m_silu_mul_quant(
const at::Tensor& input, // (E, T, 2*H)
const at::Tensor& counts, // (E)
at::Tensor& y_q, // (E, T, H) [OUT]
at::Tensor& y_s, // (E, T, H//group_size) [OUT]
bool use_ue8m0);
void gelu_and_mul(torch::Tensor& out, torch::Tensor& input);
void gelu_tanh_and_mul(torch::Tensor& out, torch::Tensor& input);
+5
View File
@@ -0,0 +1,5 @@
#include "core/registration.h"
// QuTLASS registers torch.ops._qutlass_C via TORCH_LIBRARY in bindings.cpp.
// This stub lets Python import vllm._qutlass_C to trigger op registration.
REGISTER_EXTENSION(_qutlass_C)
-40
View File
@@ -20,17 +20,6 @@
TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
// vLLM custom ops
//
ops.def(
"persistent_masked_m_silu_mul_quant(Tensor input, Tensor counts, Tensor! "
"y_q, Tensor! y_s,"
"bool use_ue8m0) -> ()");
ops.impl("persistent_masked_m_silu_mul_quant", torch::kCUDA,
&persistent_masked_m_silu_mul_quant);
ops.def("weak_ref_tensor(Tensor input) -> Tensor");
ops.impl("weak_ref_tensor", torch::kCUDA, &weak_ref_tensor);
#ifdef USE_ROCM
// TODO: Remove this once we upgrade to torch 2.11.
@@ -39,35 +28,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
ops.impl("get_cuda_view_from_cpu_tensor", torch::kCPU,
&get_cuda_view_from_cpu_tensor);
#endif
// Activation ops (quantized only — basic ops moved to _C_stable_libtorch)
ops.def(
"silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant);
// Horizontally-fused DeepseekV4-MLA: per-head RMSNorm + GPT-J RoPE for Q, and
// GPT-J RoPE + UE8M0 FP8 quant + paged cache insert for KV, all in one
// kernel launch. Registered in _C_stable_libtorch (incl. the FlashInfer V4
// full-cache bf16/fp8 variants).
// Quantization ops
#ifndef USE_ROCM
// Note about marlin kernel 'workspace' arguments:
// Technically these should be mutable since they are modified by the kernel.
// But since they are set back to zero once the kernel is finished we can
// hand wave and say that they have no net effect.
//
// The reason to mark 'workspace' as immutable is so that they don't interfere
// with using ScalarType arguments in the ops. If they are marked as mutable,
// pytorch throws an assert in
// 'torch._higher_order_ops._register_effectful_op' that prevents these
// kernels from being torch.compile'd.
// See the following document for more info on custom types and ops that use
// custom types:
// https://docs.google.com/document/d/18fBMPuOJ0fY5ZQ6YyrHUppw9FA332CpNtgB6SOIgyuA
#endif
}
+1 -1
View File
@@ -792,7 +792,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.12
ARG FLASHINFER_VERSION=0.6.13rc2
RUN --mount=type=cache,target=/opt/uv/cache \
uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
--index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
+7 -7
View File
@@ -75,13 +75,13 @@ RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRO
# Install UMD
RUN mkdir neo && \
cd neo && \
wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-core-2_2.24.8+20344_amd64.deb && \
wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-opencl-2_2.24.8+20344_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-ocloc_25.48.36300.8-0_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-opencl-icd_25.48.36300.8-0_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libigdgmm12_22.8.2_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libze-intel-gpu1_25.48.36300.8-0_amd64.deb && \
wget https://github.com/oneapi-src/level-zero/releases/download/v1.26.0/level-zero_1.26.0+u24.04_amd64.deb && \
wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.34.4/intel-igc-core-2_2.34.4+21428_amd64.deb && \
wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.34.4/intel-igc-opencl-2_2.34.4+21428_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/26.18.38308.1/intel-ocloc_26.18.38308.1-0_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/26.18.38308.1/intel-opencl-icd_26.18.38308.1-0_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/26.18.38308.1/libigdgmm12_22.10.0_amd64.deb && \
wget https://github.com/intel/compute-runtime/releases/download/26.18.38308.1/libze-intel-gpu1_26.18.38308.1-0_amd64.deb && \
wget https://github.com/oneapi-src/level-zero/releases/download/v1.28.2/level-zero_1.28.2+u24.04_amd64.deb && \
dpkg -i *.deb && \
cd .. && \
rm -rf neo
+1 -1
View File
@@ -68,7 +68,7 @@
"default": "true"
},
"FLASHINFER_VERSION": {
"default": "0.6.12"
"default": "0.6.13rc2"
},
"GDRCOPY_CUDA_VERSION": {
"default": "12.8"
+1 -1
View File
@@ -133,7 +133,7 @@ The model should inherit protocol `IsAttentionFree` and also implement class met
For the mamba layers themselves, please use the [`MambaMixer`](../../../vllm/model_executor/layers/mamba/mamba_mixer.py) (for Mamba-1) or [`MambaMixer2`](../../../vllm/model_executor/layers/mamba/mamba_mixer2.py) (for Mamba-2) classes.
The model should also be added to the `MODELS_CONFIG_MAP` dictionary in [vllm/model_executor/models/config.py](../../../vllm/model_executor/models/config.py) to ensure that the runtime defaults are optimized.
For case (2), we recommend using as a reference the implementation of [`JambaForCausalLM`](../../../vllm/model_executor/models/jamba.py) (for an example of a model that uses Mamba-1 and attention together) or [`BambaForCausalLM`](../../../vllm/model_executor/models/bamba.py) (for an example of a model that uses Mamba-2 and attention together).
For case (2), we recommend using as a reference the implementation of [`JambaForCausalLM`](../../../vllm/model_executor/models/jamba.py) (for an example of a model that uses Mamba-1 and attention together) or [`NemotronHForCausalLM`](../../../vllm/model_executor/models/nemotron_h.py) (for an example of a model that uses Mamba-2 and attention together).
These models should follow the same instructions as case (1), but they should inherit protocol `IsHybrid` (instead of `IsAttentionFree`) and it is *not* necessary to add them to the `MODELS_CONFIG_MAP` (their runtime defaults will be inferred from the protocol).
For case (3), we recommend looking at the implementation of [`MiniMaxText01ForCausalLM`](../../../vllm/model_executor/models/minimax_text_01.py) or [`Lfm2ForCausalLM`](../../../vllm/model_executor/models/lfm2.py) as a reference, which use custom "mamba-like" layers `MiniMaxText01LinearAttention` and `ShortConv` respectively.
@@ -4,7 +4,7 @@ Deploying vLLM on Kubernetes is a scalable and efficient way to serve machine le
* **Upstream vLLM compatibility** It wraps around upstream vLLM without modifying its code.
* **Ease of use** Simplified deployment via Helm charts and observability through Grafana dashboards.
* **High performance** Optimized for LLM workloads with features like multimodel support, model-aware and prefix-aware routing, fast vLLM bootstrapping, and KV cache offloading with [LMCache](https://github.com/LMCache/LMCache), among others.
* **High performance** Optimized for LLM workloads with features like multimodel support, model-aware and prefix-aware routing, fast vLLM bootstrapping, and KV cache offloading with [LMCache](https://github.com/LMCache/LMCache) (wired up in vLLM via `--kv-offloading-backend lmcache`; see the [LMCache examples](https://github.com/vllm-project/vllm/tree/main/examples/disaggregated/lmcache) and [docs.lmcache.ai](https://docs.lmcache.ai)), among others.
If you are new to Kubernetes, don't worry: in the vLLM production stack [repo](https://github.com/vllm-project/production-stack), we provide a step-by-step [guide](https://github.com/vllm-project/production-stack/blob/main/tutorials/00-install-kubernetes-env.md) and a [short video](https://www.youtube.com/watch?v=EsTJbQtzj0g) to set up everything and get started in **4 minutes**!
+1 -1
View File
@@ -20,7 +20,7 @@ Two main reasons:
Now supports 9 types of connectors:
- **ExampleConnector**: refer to [examples/disaggregated/example_connector/run.sh](../../examples/disaggregated/example_connector/run.sh) for the example usage of ExampleConnector disaggregated prefilling.
- **LMCacheConnectorV1**: refer to [examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh](../../examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh) for the example usage of LMCacheConnectorV1 disaggregated prefilling which uses NIXL as the underlying KV transmission.
- **LMCacheConnectorV1**: refer to [examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh](../../examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh) for the example usage of LMCacheConnectorV1 disaggregated prefilling which uses NIXL as the underlying KV transmission. LMCache also offers a multi-process (MP) mode via `LMCacheMPConnector`, where a standalone `lmcache server` holds the KV cache shared by one or more vLLM instances; see the [LMCache examples](../../examples/disaggregated/lmcache/README.md) and the [LMCache docs](https://docs.lmcache.ai) for setup.
- **NixlConnector**: refer to [tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh](../../tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh) for the example usage of NixlConnector disaggregated prefilling which support fully async send/recv. For detailed usage guide, see [NixlConnector Usage Guide](nixl_connector_usage.md). For feature compatibility details, see [NixlConnector Compatibility Matrix](nixl_connector_compatibility.md). You may specify one or multiple NIXL transfer backends, such as:
```bash
@@ -203,6 +203,7 @@ the vLLM JSON config.
### kv_connector_extra_config
- `load_async` (bool): Enable asynchronous loading for better compute-I/O overlap. Default: `true`.
- `lookup_async` (bool): Run the external prefix-cache lookup on a background thread so it never blocks the scheduler step. The request is held until the in-flight lookup completes, then resumed on a later step. Default: `false`.
- `enable_cross_layers_blocks` (bool): Enable cross-layer block packing for reduced store operations. Default: `false`.
- `lookup_rpc_port` (int): Custom port for the ZMQ lookup RPC socket. Default: `0`.
- `cache_prefix` (str): Namespace prepended to every store key. Lets separate deployments share one Mooncake master without polluting each other — instances configured with different prefixes never see each other's cached blocks, even for identical prompts. All instances that should share a prefix cache must use the same value. Default: `""` (no prefix; keys are byte-identical to the unprefixed format).
@@ -27,6 +27,7 @@ Currently, there are no pre-built XPU wheels.
- First, install required [driver](https://dgpu-docs.intel.com/driver/installation.html#installing-gpu-drivers).
- Second, install Python packages for vLLM XPU backend building (Intel OneAPI dependencies are installed automatically as part of `torch-xpu`, see [PyTorch XPU get started](https://docs.pytorch.org/docs/stable/notes/get_start_xpu.html)):
- Start from vllm-xpu-kernels v0.1.10, we recommend user upgrade driver to [compute runtime 26.18](https://github.com/intel/compute-runtime/releases/tag/26.14.37833.4) release, to avoid potential compatibility issue.
```bash
git clone https://github.com/vllm-project/vllm.git
+34 -14
View File
@@ -1,10 +1,38 @@
# LMCache Examples
This folder demonstrates how to use LMCache for disaggregated prefilling, CPU offloading and KV cache sharing.
This folder demonstrates how to use LMCache with vLLM v1 for KV cache
offloading, disaggregated prefilling, and KV cache sharing.
## 1. Disaggregated Prefill in vLLM v1
## Integration modes
This example demonstrates how to run LMCache with disaggregated prefill using NIXL on a single node.
LMCache integrates with vLLM v1 in two ways:
- **In-process mode** (`LMCacheConnectorV1`): LMCache runs inside the vLLM
process and is configured through environment variables or a YAML config
file (`LMCACHE_CONFIG_FILE`). This is the simplest way to add single-node
CPU/disk offloading.
- **Multi-process (MP) mode** (`LMCacheMPConnector`): LMCache runs as a
standalone server (`lmcache server`) that owns the KV cache storage; one or
more vLLM instances connect to it. This is the recommended mode for
distributed KV storage and for sharing KV cache across instances. See the
[LMCache docs](https://docs.lmcache.ai) for the full MP setup.
## 1. CPU offload (in-process)
- `python cpu_offload_lmcache.py` - CPU offloading with `LMCacheConnectorV1`
for vLLM v1.
## 2. CPU offload (multi-process)
- `bash cpu_offload_lmcache_mp.sh` - CPU offloading with `LMCacheMPConnector`,
using a standalone `lmcache server`. vLLM provides a built-in shortcut for
this setup via `--kv-offloading-backend lmcache` and
`--kv-offloading-size <GiB>`.
## 3. Disaggregated Prefill in vLLM v1
This example demonstrates how to run LMCache with disaggregated prefill using
NIXL on a single node.
### Prerequisites
@@ -46,15 +74,7 @@ The main script generates several log files:
- `decoder.log` - Logs from the decode server
- `proxy.log` - Logs from the proxy server
## 2. CPU Offload Examples
## 4. KV Cache Sharing
- `python cpu_offload_lmcache.py -v v0` - CPU offloading implementation for vLLM v0
- `python cpu_offload_lmcache.py -v v1` - CPU offloading implementation for vLLM v1
## 3. KV Cache Sharing
The `kv_cache_sharing_lmcache_v1.py` example demonstrates how to share KV caches between vLLM v1 instances.
## 4. Disaggregated Prefill in vLLM v0
The `disaggregated_prefill_lmcache_v0.py` provides an example of how to run disaggregated prefill in vLLM v0.
The `kv_cache_sharing_lmcache_v1.py` example demonstrates how to share KV
caches between vLLM v1 instances through a centralized LMCache server.
@@ -1,20 +1,8 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
This file demonstrates the example usage of cpu offloading
with LMCache in vLLM v1 or v0.
Usage:
Specify vLLM version
-v v0 : Use LMCacheConnector
model = mistralai/Mistral-7B-Instruct-v0.2
(Includes enable_chunked_prefill = True)
-v v1 : Use LMCacheConnectorV1 (default)
model = meta-llama/Meta-Llama-3.1-8B-Instruct
(Without enable_chunked_prefill)
This file demonstrates the example usage of CPU offloading
with LMCache in vLLM v1.
Note that `lmcache` is needed to run this example.
Requirements:
@@ -23,7 +11,6 @@ Learn more about LMCache environment setup, please refer to:
https://docs.lmcache.ai/getting_started/installation.html
"""
import argparse
import contextlib
import os
import time
@@ -39,8 +26,6 @@ from vllm.engine.arg_utils import EngineArgs
def setup_environment_variables():
# LMCache-related environment variables
# Use experimental features in LMCache
os.environ["LMCACHE_USE_EXPERIMENTAL"] = "True"
# LMCache is set to use 256 tokens per chunk
os.environ["LMCACHE_CHUNK_SIZE"] = "256"
# Enable local CPU backend in LMCache
@@ -50,9 +35,9 @@ def setup_environment_variables():
@contextlib.contextmanager
def build_llm_with_lmcache(lmcache_connector: str, model: str):
def build_llm_with_lmcache(model: str):
ktc = KVTransferConfig(
kv_connector=lmcache_connector,
kv_connector="LMCacheConnectorV1",
kv_role="kv_both",
)
# Set GPU memory utilization to 0.8 for an A40 GPU with 40GB
@@ -92,23 +77,10 @@ def print_output(
print("-" * 50)
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument(
"-v",
"--version",
choices=["v0", "v1"],
default="v1",
help="Specify vLLM version (default: v1)",
)
return parser.parse_args()
def main():
lmcache_connector = "LMCacheConnectorV1"
model = "meta-llama/Meta-Llama-3.1-8B-Instruct"
setup_environment_variables()
with build_llm_with_lmcache(lmcache_connector, model) as llm:
with build_llm_with_lmcache(model) as llm:
# This example script runs two requests with a shared prefix.
# Define the shared prompt and specific prompts
shared_prompt = "Hello, how are you?" * 1000
+43
View File
@@ -0,0 +1,43 @@
#!/bin/bash
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
#
# CPU offloading with LMCache in multi-process (MP) mode.
#
# In MP mode, LMCache runs as a standalone server process (`lmcache server`)
# that owns the KV cache storage. One or more vLLM instances connect to it via
# the `LMCacheMPConnector`. This is the recommended way to run LMCache for
# distributed KV storage and for sharing KV cache across vLLM instances.
#
# vLLM ships a built-in shortcut for this setup: pass `--kv-offloading-backend
# lmcache` together with `--kv-offloading-size <GiB>` and vLLM wires up the
# `LMCacheMPConnector` for you (it defaults to the LMCache server at
# tcp://localhost:5555, matching the `lmcache server` default).
#
# Requires `lmcache` to be installed (`pip install lmcache`).
# Learn more: https://docs.lmcache.ai
set -euo pipefail
MODEL=${MODEL:-meta-llama/Meta-Llama-3.1-8B-Instruct}
# 1. Launch the standalone LMCache server (binds tcp://localhost:5555 by
# default). `--l1-size-gb` sets the CPU memory budget for the L1 cache.
echo "Starting LMCache server..."
lmcache server --host localhost --port 5555 --l1-size-gb 5 &
LMCACHE_SERVER_PID=$!
trap 'kill $LMCACHE_SERVER_PID 2>/dev/null || true' EXIT
# 2. Launch vLLM and offload KV cache to the LMCache server.
# The MP connector currently requires the non-hybrid KV cache manager.
echo "Starting vLLM server with LMCache MP offloading..."
vllm serve "$MODEL" \
--port 8000 \
--kv-offloading-size 5 \
--kv-offloading-backend lmcache \
--disable-hybrid-kv-cache-manager
# Equivalent explicit configuration (instead of the two flags above):
# --kv-transfer-config \
# '{"kv_connector":"LMCacheMPConnector","kv_role":"kv_both",
# "kv_connector_extra_config":{"lmcache.mp.host":"tcp://localhost",
# "lmcache.mp.port":5555}}'
@@ -1,144 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
This file demonstrates the example usage of disaggregated prefilling
with LMCache.
We will launch 2 vllm instances (GPU 0 for prefill and GPU 1 for decode),
and launch an additional LMCache server.
KV cache is transferred in the following manner:
vLLM prefill node -> LMCache server -> vLLM decode node.
Note that `pip install lmcache` is needed to run this example.
Learn more about LMCache in https://github.com/LMCache/LMCache.
"""
import os
import subprocess
import time
from multiprocessing import Event, Process
from lmcache.experimental.cache_engine import LMCacheEngineBuilder
from lmcache.integration.vllm.utils import ENGINE_NAME
from vllm import LLM, SamplingParams
from vllm.config import KVTransferConfig
# LMCache-related environment variables
# The port to start LMCache server
port = 8100
# Use experimental features in LMCache
os.environ["LMCACHE_USE_EXPERIMENTAL"] = "True"
# LMCache is set to use 256 tokens per chunk
os.environ["LMCACHE_CHUNK_SIZE"] = "256"
# Disable local CPU backend in LMCache
os.environ["LMCACHE_LOCAL_CPU"] = "False"
# Set local CPU memory buffer limit to 5.0 GB
os.environ["LMCACHE_MAX_LOCAL_CPU_SIZE"] = "5.0"
# Set the remote URL for LMCache server
os.environ["LMCACHE_REMOTE_URL"] = f"lm://localhost:{port}"
# Set the serializer/deserializer between vllm and LMCache server
# `naive` indicates using raw bytes of the tensor without any compression
os.environ["LMCACHE_REMOTE_SERDE"] = "naive"
prompts = [
"Hello, how are you?" * 1000,
]
def run_prefill(prefill_done, prompts):
# We use GPU 0 for prefill node.
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=1)
ktc = KVTransferConfig(
kv_connector="LMCacheConnector",
kv_role="kv_producer",
kv_rank=0,
kv_parallel_size=2,
)
# Set GPU memory utilization to 0.8 for an A40 GPU with 40GB
# memory. Reduce the value if your GPU has less memory.
llm = LLM(
model="mistralai/Mistral-7B-Instruct-v0.2",
kv_transfer_config=ktc,
max_model_len=8000,
gpu_memory_utilization=0.8,
enforce_eager=True,
)
# llm.generate(prompts, sampling_params)
outputs = llm.generate(prompts, sampling_params)
for output in outputs:
generated_text = output.outputs[0].text
print(f"Generated text: {generated_text!r}")
print("Prefill node is finished.")
prefill_done.set()
# Clean up lmcache backend
LMCacheEngineBuilder.destroy(ENGINE_NAME)
def run_decode(prefill_done, prompts, timeout=1):
# We use GPU 1 for decode node.
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10)
ktc = KVTransferConfig(
kv_connector="LMCacheConnector",
kv_role="kv_consumer",
kv_rank=1,
kv_parallel_size=2,
)
# Set GPU memory utilization to 0.8 for an A40 GPU with 40GB
# of memory. Reduce the value if your GPU has less memory.
llm = LLM(
model="mistralai/Mistral-7B-Instruct-v0.2",
kv_transfer_config=ktc,
max_model_len=8000,
gpu_memory_utilization=0.8,
enforce_eager=True,
)
print("Waiting for prefill node to finish...")
prefill_done.wait()
time.sleep(timeout)
outputs = llm.generate(prompts, sampling_params)
for output in outputs:
generated_text = output.outputs[0].text
print(f"Generated text: {generated_text!r}")
# Clean up lmcache backend
LMCacheEngineBuilder.destroy(ENGINE_NAME)
def run_lmcache_server(port):
server_proc = subprocess.Popen(
["python", "-m", "lmcache.experimental.server", "localhost", str(port)]
)
return server_proc
def main():
prefill_done = Event()
prefill_process = Process(target=run_prefill, args=(prefill_done, prompts))
decode_process = Process(target=run_decode, args=(prefill_done, prompts))
lmcache_server_process = run_lmcache_server(port)
# Start prefill node
prefill_process.start()
# Start decode node
decode_process.start()
# Clean up the processes
decode_process.join()
prefill_process.terminate()
lmcache_server_process.terminate()
lmcache_server_process.wait()
if __name__ == "__main__":
main()
@@ -30,7 +30,6 @@ if [[ $1 == "prefiller" ]]; then
UCX_TLS=cuda_ipc,cuda_copy,tcp \
LMCACHE_CONFIG_FILE=$prefill_config_file \
LMCACHE_USE_EXPERIMENTAL=True \
VLLM_ENABLE_V1_MULTIPROCESSING=1 \
VLLM_WORKER_MULTIPROC_METHOD=spawn \
CUDA_VISIBLE_DEVICES=0 \
@@ -47,7 +46,6 @@ elif [[ $1 == "decoder" ]]; then
UCX_TLS=cuda_ipc,cuda_copy,tcp \
LMCACHE_CONFIG_FILE=$decode_config_file \
LMCACHE_USE_EXPERIMENTAL=True \
VLLM_ENABLE_V1_MULTIPROCESSING=1 \
VLLM_WORKER_MULTIPROC_METHOD=spawn \
CUDA_VISIBLE_DEVICES=1 \
@@ -26,8 +26,6 @@ from vllm.config import KVTransferConfig
# LMCache-related environment variables
# The port to start LMCache server
port = 8100
# Use experimental features in LMCache
os.environ["LMCACHE_USE_EXPERIMENTAL"] = "True"
# LMCache is set to use 256 tokens per chunk
os.environ["LMCACHE_CHUNK_SIZE"] = "256"
# Disable local CPU backend in LMCache
+3 -2
View File
@@ -11,13 +11,14 @@ transformers >= 5.5.3
tokenizers >= 0.21.1 # Required for fast incremental detokenization.
safetensors >= 0.6.2 # MXFP4/MXFP6 dtype support (F8_E8M0, F4) added in 0.6.0: https://github.com/huggingface/safetensors/pull/611
protobuf >= 5.29.6, !=6.30.*, !=6.31.*, !=6.32.*, !=6.33.0.*, !=6.33.1.*, !=6.33.2.*, !=6.33.3.*, !=6.33.4.* # Required by LlamaTokenizer, gRPC. CVE-2026-0994
fastapi[standard] >= 0.115.0 # Required by FastAPI's form models in the OpenAI API server's audio transcriptions endpoint.
fastapi[standard] >= 0.133.0, < 0.137.0 # First version supporting Starlette 1.0; < 0.137.0 avoids route-tree change that breaks model-hosting-container-standards handler overrides.
starlette >= 1.0.1 # CVE-2026-48710: Host header injection in < 1.0.1
aiohttp >= 3.13.3
openai >= 2.0.0 # For Responses API with reasoning content
pydantic >= 2.12.0
prometheus_client >= 0.18.0
pillow # Required for image processing
prometheus-fastapi-instrumentator >= 7.0.0
prometheus-fastapi-instrumentator >= 8.0.0 # v8 unblocks starlette >= 1.0
tiktoken >= 0.6.0 # Required for DBRX tokenizer
lm-format-enforcer == 0.11.3
llguidance >= 1.7.0, < 1.8.0; platform_machine == "x86_64" or platform_machine == "arm64" or platform_machine == "aarch64" or platform_machine == "ppc64le"
+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.12
flashinfer-cubin==0.6.12
flashinfer-python==0.6.13rc2
flashinfer-cubin==0.6.13rc2
apache-tvm-ffi==0.1.9
tilelang==0.1.9
nvidia-cudnn-frontend>=1.19.1
+1 -1
View File
@@ -40,7 +40,7 @@ lm-eval[api]>=0.4.12 # required for model evaluation test
mteb[bm25s]>=2, <3 # required for mteb test
transformers==5.5.3
tokenizers==0.22.2
schemathesis>=3.39.15 # Required for openai schema test.
schemathesis>=4.0.0 # Required for openai schema test.
# quantization
bitsandbytes==0.49.2
buildkite-test-collector==0.1.9
+23 -48
View File
@@ -35,14 +35,11 @@ arctic-inference==0.1.1
# via -r requirements/test/cuda.in
argcomplete==3.5.1
# via datamodel-code-generator
arrow==1.3.0
# via isoduration
attrs==24.2.0
# via
# aiohttp
# hypothesis
# jsonschema
# pytest-subtests
# referencing
audioread==3.0.1
# via librosa
@@ -57,9 +54,7 @@ azure-identity==1.25.2
azure-storage-blob==12.28.0
# via runai-model-streamer-azure
backoff==2.2.1
# via
# -r requirements/test/cuda.in
# schemathesis
# via -r requirements/test/cuda.in
bitsandbytes==0.49.2
# via -r requirements/test/cuda.in
black==24.10.0
@@ -110,7 +105,6 @@ colorama==0.4.6
# via
# perceptron
# sacrebleu
# schemathesis
colorful==0.5.6
# via ray
colorlog==6.10.1
@@ -183,7 +177,7 @@ et-xmlfile==2.0.0
# via openpyxl
evaluate==0.4.3
# via lm-eval
fastapi==0.128.0
fastapi==0.136.3
# via
# -c requirements/common.txt
# gpt-oss
@@ -206,8 +200,6 @@ filelock==3.16.1
# virtualenv
fonttools==4.55.0
# via matplotlib
fqdn==1.5.1
# via jsonschema
frozendict==2.4.6
# via einx
frozenlist==1.5.0
@@ -269,7 +261,7 @@ h11==0.14.0
# uvicorn
h2==4.3.0
# via httpx
harfile==0.3.0
harfile==0.5.0
# via schemathesis
hf-xet==1.4.3
# via huggingface-hub
@@ -309,7 +301,7 @@ hypothesis==6.131.0
# hypothesis-graphql
# hypothesis-jsonschema
# schemathesis
hypothesis-graphql==0.11.1
hypothesis-graphql==0.13.0
# via schemathesis
hypothesis-jsonschema==0.23.1
# via schemathesis
@@ -318,7 +310,6 @@ idna==3.10
# anyio
# email-validator
# httpx
# jsonschema
# requests
# yarl
imagehash==4.3.2
@@ -335,8 +326,6 @@ instanttensor==0.1.5
# via -r requirements/test/cuda.in
isodate==0.7.2
# via azure-storage-blob
isoduration==20.11.0
# via jsonschema
isort==5.13.2
# via datamodel-code-generator
jinja2==3.1.6
@@ -356,15 +345,14 @@ joblib==1.4.2
# librosa
# nltk
# scikit-learn
jsonpointer==3.0.0
# via jsonschema
jsonschema==4.23.0
# via
# -c requirements/common.txt
# hypothesis-jsonschema
# mistral-common
# ray
# schemathesis
jsonschema-rs==0.46.5
# via schemathesis
jsonschema-specifications==2024.10.1
# via jsonschema
junit-xml==1.9
@@ -715,18 +703,20 @@ pydantic-core==2.41.1
pydantic-extra-types==2.10.5
# via mistral-common
pygments==2.18.0
# via rich
# via
# pytest
# rich
pyjwt==2.11.0
# via msal
pyparsing==3.2.0
# via matplotlib
pyrate-limiter==3.7.0
pyrate-limiter==4.4.0
# via schemathesis
pystemmer==3.0.0
# via mteb
pytablewriter==1.2.0
# via lm-eval
pytest==8.3.5
pytest==9.1.0
# via
# -r requirements/test/cuda.in
# buildkite-test-collector
@@ -737,10 +727,9 @@ pytest==8.3.5
# pytest-mock
# pytest-rerunfailures
# pytest-shard
# pytest-subtests
# pytest-timeout
# schemathesis
pytest-asyncio==0.24.0
pytest-asyncio==1.4.0
# via -r requirements/test/cuda.in
pytest-cov==6.3.0
# via -r requirements/test/cuda.in
@@ -752,13 +741,10 @@ pytest-rerunfailures==14.0
# via -r requirements/test/cuda.in
pytest-shard==0.1.2
# via -r requirements/test/cuda.in
pytest-subtests==0.14.1
# via schemathesis
pytest-timeout==2.3.1
# via -r requirements/test/cuda.in
python-dateutil==2.9.0.post0
# via
# arrow
# botocore
# matplotlib
# pandas
@@ -829,15 +815,12 @@ requests==2.32.3
# tiktoken
responses==0.25.3
# via genai-perf
rfc3339-validator==0.1.4
# via jsonschema
rfc3987==1.3.8
# via jsonschema
rich==13.9.4
# via
# genai-perf
# mteb
# perceptron
# schemathesis
# typer
rouge-score==0.1.2
# via lm-eval
@@ -868,7 +851,7 @@ safetensors==0.7.0
# segmentation-models-pytorch
# timm
# transformers
schemathesis==3.39.15
schemathesis==4.21.6
# via -r requirements/test/cuda.in
scikit-image==0.25.2
# via albumentations
@@ -912,7 +895,6 @@ six==1.16.0
# junit-xml
# opencensus
# python-dateutil
# rfc3339-validator
# rouge-score
smart-open==7.1.0
# via ray
@@ -938,10 +920,10 @@ sqlalchemy==2.0.41
# optuna
sqlitedict==2.1.0
# via lm-eval
starlette==0.50.0
starlette==1.3.1
# via
# -c requirements/common.txt
# fastapi
# schemathesis
# starlette-testclient
starlette-testclient==0.4.1
# via schemathesis
@@ -966,6 +948,7 @@ tenacity==9.1.2
# gpt-oss
# lm-eval
# plotly
# schemathesis
tensorizer==2.10.1
# via -r requirements/test/cuda.in
termcolor==3.1.0
@@ -990,10 +973,6 @@ tokenizers==0.22.2
# -c requirements/common.txt
# -r requirements/test/cuda.in
# transformers
tomli==2.2.1
# via schemathesis
tomli-w==1.2.0
# via schemathesis
torch==2.11.0+cu130
# via
# -c requirements/cuda.txt
@@ -1066,8 +1045,6 @@ typer==0.15.2
# huggingface-hub
# perceptron
# transformers
types-python-dateutil==2.9.0.20241206
# via arrow
typing-extensions==4.15.0
# via
# -c requirements/common.txt
@@ -1092,6 +1069,8 @@ typing-extensions==4.15.0
# pydantic
# pydantic-core
# pydantic-extra-types
# pytest-asyncio
# schemathesis
# sentence-transformers
# sqlalchemy
# starlette
@@ -1099,11 +1078,11 @@ typing-extensions==4.15.0
# typer
# typing-inspection
typing-inspection==0.4.2
# via pydantic
# via
# fastapi
# pydantic
tzdata==2024.2
# via pandas
uri-template==1.3.0
# via jsonschema
urllib3==2.2.3
# via
# blobfile
@@ -1122,8 +1101,6 @@ vocos==0.1.0
# via -r requirements/test/cuda.in
wcwidth==0.2.13
# via ftfy
webcolors==24.11.1
# via jsonschema
werkzeug==3.1.3
# via schemathesis
word2number==1.1
@@ -1135,8 +1112,6 @@ xxhash==3.5.0
# datasets
# evaluate
yarl==1.17.1
# via
# aiohttp
# schemathesis
# via aiohttp
zipp==3.23.0
# via importlib-metadata
+1 -1
View File
@@ -31,7 +31,7 @@ lm-eval[api]>=0.4.12 # required for model evaluation test
mteb[bm25s]>=2, <3 # required for mteb test
transformers==5.5.3
tokenizers==0.22.2
schemathesis>=3.39.15 # Required for openai schema test.
schemathesis>=4.0.0 # Required for openai schema test.
# quantization
bitsandbytes>=0.49.2
buildkite-test-collector==0.1.9
+1 -1
View File
@@ -39,7 +39,7 @@ lm-eval[api]>=0.4.12 # required for model evaluation test
mteb[bm25s]>=2, <3 # required for mteb test
transformers==5.5.3
tokenizers==0.22.2
schemathesis>=3.39.15 # Required for openai schema test
schemathesis>=4.0.0 # Required for openai schema test
# quantization
bitsandbytes==0.49.2
buildkite-test-collector==0.1.9
+21 -47
View File
@@ -51,15 +51,12 @@ arctic-inference==0.1.1
# via -r requirements/test/rocm.in
argcomplete==3.6.3
# via datamodel-code-generator
arrow==1.4.0
# via isoduration
astor==0.8.1
# via depyf
attrs==26.1.0
# via
# aiohttp
# jsonschema
# pytest-subtests
# referencing
audioread==3.0.1
# via librosa
@@ -74,9 +71,7 @@ azure-identity==1.25.3
azure-storage-blob==12.28.0
# via runai-model-streamer-azure
backoff==2.2.1
# via
# -r requirements/test/rocm.in
# schemathesis
# via -r requirements/test/rocm.in
bitsandbytes==0.49.2
# via -r requirements/test/rocm.in
black==26.3.1
@@ -139,7 +134,6 @@ colorama==0.4.6
# via
# perceptron
# sacrebleu
# schemathesis
colorful==0.5.8
# via ray
colorlog==6.10.1
@@ -258,8 +252,6 @@ filelock==3.25.2
# virtualenv
fonttools==4.62.1
# via matplotlib
fqdn==1.5.1
# via jsonschema
frozendict==2.4.7
# via einx
frozenlist==1.8.0
@@ -328,7 +320,7 @@ h11==0.16.0
# uvicorn
h2==4.3.0
# via httpx
harfile==0.4.0
harfile==0.5.0
# via schemathesis
hf-xet==1.4.3
# via huggingface-hub
@@ -378,7 +370,7 @@ hypothesis==6.151.9
# hypothesis-graphql
# hypothesis-jsonschema
# schemathesis
hypothesis-graphql==0.12.0
hypothesis-graphql==0.13.0
# via schemathesis
hypothesis-jsonschema==0.23.1
# via schemathesis
@@ -387,7 +379,6 @@ idna==3.11
# anyio
# email-validator
# httpx
# jsonschema
# requests
# yarl
ijson==3.5.0
@@ -408,8 +399,6 @@ interegular==0.3.3
# via lm-format-enforcer
isodate==0.7.2
# via azure-storage-blob
isoduration==20.11.0
# via jsonschema
isort==8.0.1
# via datamodel-code-generator
jinja2==3.1.6
@@ -435,8 +424,6 @@ joblib==1.5.3
# librosa
# nltk
# scikit-learn
jsonpointer==3.1.0
# via jsonschema
jsonschema==4.26.0
# via
# -c requirements/common.txt
@@ -445,7 +432,8 @@ jsonschema==4.26.0
# mcp
# mistral-common
# ray
# schemathesis
jsonschema-rs==0.46.5
# via schemathesis
jsonschema-specifications==2025.9.1
# via jsonschema
junit-xml==1.9
@@ -792,7 +780,7 @@ prometheus-client==0.24.1
# opentelemetry-exporter-prometheus
# prometheus-fastapi-instrumentator
# ray
prometheus-fastapi-instrumentator==7.1.0
prometheus-fastapi-instrumentator==8.0.0
# via
# -c requirements/common.txt
# -r requirements/test/../common.txt
@@ -876,20 +864,22 @@ pydantic-settings==2.13.1
# fastapi
# mcp
pygments==2.19.2
# via rich
# via
# pytest
# rich
pyjwt==2.12.1
# via
# mcp
# msal
pyparsing==3.3.2
# via matplotlib
pyrate-limiter==3.9.0
pyrate-limiter==4.4.0
# via schemathesis
pystemmer==3.0.0
# via mteb
pytablewriter==1.2.1
# via lm-eval
pytest==8.3.5
pytest==9.1.0
# via
# -r requirements/test/rocm.in
# buildkite-test-collector
@@ -900,10 +890,9 @@ pytest==8.3.5
# pytest-mock
# pytest-rerunfailures
# pytest-shard
# pytest-subtests
# pytest-timeout
# schemathesis
pytest-asyncio==0.24.0
pytest-asyncio==1.4.0
# via -r requirements/test/rocm.in
pytest-cov==6.3.0
# via -r requirements/test/rocm.in
@@ -915,13 +904,10 @@ pytest-rerunfailures==14.0
# via -r requirements/test/rocm.in
pytest-shard==0.1.2
# via -r requirements/test/rocm.in
pytest-subtests==0.14.2
# via schemathesis
pytest-timeout==2.3.1
# via -r requirements/test/rocm.in
python-dateutil==2.9.0.post0
# via
# arrow
# botocore
# matplotlib
# pandas
@@ -1016,16 +1002,13 @@ requests==2.32.5
# tiktoken
responses==0.26.0
# via genai-perf
rfc3339-validator==0.1.4
# via jsonschema
rfc3987==1.3.8
# via jsonschema
rich==14.3.3
# via
# genai-perf
# mteb
# perceptron
# rich-toolkit
# schemathesis
# typer
rich-toolkit==0.19.7
# via
@@ -1063,7 +1046,7 @@ safetensors==0.7.0
# segmentation-models-pytorch
# timm
# transformers
schemathesis==3.39.15
schemathesis==4.21.6
# via -r requirements/test/rocm.in
scikit-image==0.26.0
# via albumentations
@@ -1120,7 +1103,6 @@ six==1.17.0
# junit-xml
# opencensus
# python-dateutil
# rfc3339-validator
# rouge-score
smart-open==7.5.1
# via ray
@@ -1149,13 +1131,14 @@ sqlitedict==2.1.0
# via lm-eval
sse-starlette==3.3.4
# via mcp
starlette==0.52.1
starlette==1.3.1
# via
# -c requirements/common.txt
# -r requirements/test/../common.txt
# fastapi
# mcp
# model-hosting-container-standards
# prometheus-fastapi-instrumentator
# schemathesis
# sse-starlette
# starlette-testclient
starlette-testclient==0.4.1
@@ -1182,6 +1165,7 @@ tenacity==9.1.4
# via
# gpt-oss
# lm-eval
# schemathesis
tensorizer==2.10.1
# via
# -c requirements/rocm.txt
@@ -1215,10 +1199,6 @@ tokenizers==0.22.2
# -r requirements/test/../common.txt
# -r requirements/test/rocm.in
# transformers
tomli==2.4.0
# via schemathesis
tomli-w==1.2.0
# via schemathesis
torch-c-dlpack-ext==0.1.5
# via tilelang
tqdm==4.67.3
@@ -1301,8 +1281,10 @@ typing-extensions==4.15.0
# pydantic
# pydantic-core
# pydantic-extra-types
# pytest-asyncio
# referencing
# rich-toolkit
# schemathesis
# sentence-transformers
# sqlalchemy
# starlette
@@ -1317,10 +1299,6 @@ typing-inspection==0.4.2
# mcp
# pydantic
# pydantic-settings
tzdata==2025.3
# via arrow
uri-template==1.3.0
# via jsonschema
urllib3==2.6.3
# via
# blobfile
@@ -1351,8 +1329,6 @@ watchfiles==1.1.1
# uvicorn
wcwidth==0.6.0
# via ftfy
webcolors==25.10.0
# via jsonschema
websockets==16.0
# via uvicorn
werkzeug==3.1.6
@@ -1370,9 +1346,7 @@ xxhash==3.6.0
# datasets
# evaluate
yarl==1.23.0
# via
# aiohttp
# schemathesis
# via aiohttp
z3-solver==4.15.4.0
# via tilelang
zipp==3.23.0
+2 -1
View File
@@ -593,8 +593,9 @@ soxr==0.5.0.post1
# mistral-common
sqlitedict==2.1.0
# via lm-eval
starlette==1.0.0
starlette==1.3.1
# via
# -c requirements/common.txt
# fastapi
# starlette-testclient
starlette-testclient==0.4.1
+1 -1
View File
@@ -17,4 +17,4 @@ torchaudio
torchvision
auto_round_lib>=0.13.3
vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.9.1/vllm_xpu_kernels-0.1.9.1-cp38-abi3-manylinux_2_28_x86_64.whl
vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.10/vllm_xpu_kernels-0.1.10-cp38-abi3-manylinux_2_28_x86_64.whl
+12 -1
View File
@@ -769,6 +769,7 @@ class precompiled_wheel_utils:
"vllm/_C.abi3.so",
"vllm/_C_stable_libtorch.abi3.so",
"vllm/_moe_C_stable_libtorch.abi3.so",
"vllm/_qutlass_C.abi3.so",
"vllm/_flashmla_C.abi3.so",
"vllm/_flashmla_extension_C.abi3.so",
"vllm/_sparse_flashmla_C.abi3.so",
@@ -1135,6 +1136,7 @@ if _is_cuda():
# DeepGEMM requires CUDA 12.3+ (SM90/SM100)
# Optional since it won't build on unsupported architectures
ext_modules.append(CMakeExtension(name="vllm._deep_gemm_C", optional=True))
ext_modules.append(CMakeExtension(name="vllm._qutlass_C", optional=True))
# fmha_sm100 is a Python/CuTe-DSL package installed into vllm.third_party.
ext_modules.append(CMakeExtension(name="vllm.fmha_sm100", optional=True))
@@ -1149,7 +1151,8 @@ if _is_cpu():
ext_modules.append(CMakeExtension(name="vllm._C"))
if _build_custom_ops():
ext_modules.append(CMakeExtension(name="vllm._C"))
if _is_hip():
ext_modules.append(CMakeExtension(name="vllm._C"))
if _is_cuda() or _is_hip():
ext_modules.append(CMakeExtension(name="vllm._C_stable_libtorch"))
ext_modules.append(CMakeExtension(name="vllm._moe_C_stable_libtorch"))
@@ -1168,7 +1171,15 @@ package_data = {
"third_party/deep_gemm/include/**/*.h",
"third_party/deep_gemm/include/**/*.hpp",
# fmha_sm100 sparse CuTe-DSL helper kernels (vendored via cmake)
"third_party/fmha_sm100/csrc/**/*.cu",
"third_party/fmha_sm100/csrc/**/*.h",
"third_party/fmha_sm100/csrc/**/*.jinja",
"third_party/fmha_sm100/csrc/**/*.cu.jinja",
"third_party/fmha_sm100/cute/**/*.cu",
"third_party/fmha_sm100/cutlass/include/**/*.h",
"third_party/fmha_sm100/cutlass/include/**/*.hpp",
"third_party/fmha_sm100/cutlass/tools/util/include/**/*.h",
"third_party/fmha_sm100/cutlass/tools/util/include/**/*.hpp",
]
}
+193
View File
@@ -649,3 +649,196 @@ def test_cloud_storage_tokenizer_skips_get_model_path(monkeypatch):
args = EngineArgs(model="s3://bucket/model", tokenizer="s3://bucket/tokenizer")
assert args.model == "s3://bucket/model"
assert args.tokenizer == "s3://bucket/tokenizer"
class TestDeviceIds:
def test_device_ids_with_cvd_out_of_range(self, monkeypatch):
"""--device-ids index beyond the CVD set raises ValueError."""
from vllm.platforms import current_platform
key = current_platform.device_control_env_var
monkeypatch.setenv(key, "4,5")
args = EngineArgs(model="m", device_ids=[0, 2])
with pytest.raises(ValueError, match="out of range"):
args._resolve_device_ids()
def test_device_ids_with_cvd_resolve_to_physical_ids(self, monkeypatch):
"""--device-ids are CVD-local indices resolved to physical ids."""
from vllm.platforms import current_platform
key = current_platform.device_control_env_var
monkeypatch.setenv(key, "4,5")
args = EngineArgs(model="m", device_ids=[0, 1])
assert args._resolve_device_ids() == [4, 5]
def test_device_ids_with_uuid_cvd_resolve_to_physical_ids(self, monkeypatch):
"""--device-ids support UUID CVD values resolved by the platform."""
from vllm.platforms import current_platform
key = current_platform.device_control_env_var
monkeypatch.setenv(key, "GPU-abcd1234,GPU-ef567890")
monkeypatch.setattr(
type(current_platform),
"device_control_id_to_physical_device_id",
classmethod(
lambda cls, device_id: {"GPU-abcd1234": 4, "GPU-ef567890": 5}[device_id]
),
)
args = EngineArgs(model="m", device_ids=[0, 1])
assert args._resolve_device_ids() == [4, 5]
def test_device_ids_with_uuid_args_resolve_to_physical_ids(self, monkeypatch):
"""UUID --device-ids are resolved to physical IDs immediately."""
from vllm.platforms import current_platform
monkeypatch.setattr(
type(current_platform),
"device_control_id_to_physical_device_id",
classmethod(lambda cls, device_id: {"GPU-abcd1234": 4}[device_id]),
)
args = EngineArgs(model="m", device_ids=["GPU-abcd1234"])
assert args._resolve_device_ids() == [4]
def test_device_ids_reject_mixed_integer_and_uuid_args(self):
"""--device-ids must not mix CVD indices and UUIDs."""
args = EngineArgs(model="m", device_ids=[0, "GPU-abcd1234"])
with pytest.raises(ValueError, match="must not mix"):
args._resolve_device_ids()
def test_no_device_ids(self):
"""No --device-ids returns None."""
args = EngineArgs(model="m")
assert args._resolve_device_ids() is None
def test_cli_parsing(self):
"""--device-ids parses comma-separated string from CLI."""
parser = FlexibleArgumentParser()
EngineArgs.add_cli_args(parser)
parsed = parser.parse_args(["--model", "m", "--device-ids", "0,2,4"])
assert parsed.device_ids == [0, 2, 4]
def test_cli_parsing_uuid(self):
"""--device-ids parses comma-separated UUID strings from CLI."""
parser = FlexibleArgumentParser()
EngineArgs.add_cli_args(parser)
parsed = parser.parse_args(
["--model", "m", "--device-ids", "GPU-abcd1234,GPU-ef567890"]
)
assert parsed.device_ids == ["GPU-abcd1234", "GPU-ef567890"]
def test_assigned_physical_gpu_ids_are_physical_with_cvd(self, monkeypatch):
"""assigned_physical_gpu_ids are already physical and not composed with CVD."""
import vllm.platforms.interface as platform_interface
from vllm.platforms import current_platform
monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", [4, 5])
monkeypatch.setenv(current_platform.device_control_env_var, "4,5")
assert current_platform.device_id_to_physical_device_id(0) == 4
assert current_platform.device_id_to_physical_device_id(1) == 5
assert current_platform.logical_device_id_to_visible_device_id(0) == 0
assert current_platform.logical_device_id_to_visible_device_id(1) == 1
def test_assigned_physical_gpu_ids_map_to_visible_uuid_cvd(self, monkeypatch):
"""Physical IDs map back to visible ordinals when CVD uses UUIDs."""
import vllm.platforms.interface as platform_interface
from vllm.platforms import current_platform
monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", [5])
monkeypatch.setenv(
current_platform.device_control_env_var,
"GPU-abcd1234,GPU-ef567890",
)
monkeypatch.setattr(
type(current_platform),
"device_control_id_to_physical_device_id",
classmethod(
lambda cls, device_id: {"GPU-abcd1234": 4, "GPU-ef567890": 5}[device_id]
),
)
assert current_platform.logical_device_id_to_visible_device_id(0) == 1
def test_device_ids_reject_duplicates(self):
"""--device-ids must not contain duplicate entries."""
args = EngineArgs(model="m", device_ids=[2, 2])
with pytest.raises(ValueError, match="duplicates"):
args._resolve_device_ids()
def test_cli_parsing_strips_whitespace(self):
"""--device-ids tolerates whitespace around commas."""
parser = FlexibleArgumentParser()
EngineArgs.add_cli_args(parser)
parsed = parser.parse_args(["--model", "m", "--device-ids", "0, 2, 4"])
assert parsed.device_ids == [0, 2, 4]
def test_visible_ordinal_to_physical_ignores_assigned_ids(self, monkeypatch):
"""visible_device_id_to_physical_device_id maps torch device ordinals,
independent of the logical-to-physical mapping.
Regression test: CustomAllreduce passes device.index (a visible
ordinal) and must not index into assigned_physical_gpu_ids, which
raised IndexError for non-identity --device-ids like [2, 3].
"""
import vllm.platforms.interface as platform_interface
from vllm.platforms import current_platform
monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", [2, 3])
monkeypatch.delenv(current_platform.device_control_env_var, raising=False)
# CVD unset: visible ordinal == physical ID, even beyond the
# assigned list's length.
assert current_platform.visible_device_id_to_physical_device_id(2) == 2
assert current_platform.visible_device_id_to_physical_device_id(3) == 3
monkeypatch.setenv(current_platform.device_control_env_var, "4,5")
assert current_platform.visible_device_id_to_physical_device_id(1) == 5
with pytest.raises(IndexError, match="out of range"):
current_platform.visible_device_id_to_physical_device_id(2)
class TestDpDeviceIdSharding:
def test_dp_supervisor_device_ids_stay_env_relative(self):
"""Regression test: the DP supervisor must pass env-relative indices,
not physical IDs, because each child re-resolves --device-ids
against its inherited device-control env var."""
import argparse
from vllm.entrypoints.openai.dp_supervisor import _build_device_ids
args = argparse.Namespace(
tensor_parallel_size=2, pipeline_parallel_size=1, device_ids=None
)
assert _build_device_ids(args, local_rank=0) == [0, 1]
assert _build_device_ids(args, local_rank=1) == [2, 3]
def test_dp_supervisor_shards_user_device_ids(self):
"""User-provided --device-ids are sharded across DP children."""
import argparse
from vllm.entrypoints.openai.dp_supervisor import _build_device_ids
args = argparse.Namespace(
tensor_parallel_size=2, pipeline_parallel_size=1, device_ids=[4, 5, 6, 7]
)
assert _build_device_ids(args, local_rank=0) == [4, 5]
assert _build_device_ids(args, local_rank=1) == [6, 7]
with pytest.raises(ValueError, match="needs devices"):
_build_device_ids(args, local_rank=2)
def test_dp_rank_shards_user_assigned_gpu_ids(self):
"""get_physical_gpu_ids_for_local_dp_rank slices the user-provided
--device-ids list instead of recomputing from the env var."""
from vllm.platforms import current_platform
from vllm.v1.engine.utils import get_physical_gpu_ids_for_local_dp_rank
evar = current_platform.device_control_env_var
assert get_physical_gpu_ids_for_local_dp_rank(
evar, local_dp_rank=1, world_size=2, user_assigned_gpu_ids=[4, 5, 6, 7]
) == [6, 7]
with pytest.raises(ValueError, match="needs devices"):
get_physical_gpu_ids_for_local_dp_rank(
evar, local_dp_rank=2, world_size=2, user_assigned_gpu_ids=[4, 5, 6, 7]
)
@@ -1096,3 +1096,50 @@ class TestMessageStartIncludesTypeAndRole:
message = events[0][1]["message"]
assert message["type"] == "message"
assert message["role"] == "assistant"
# ======================================================================
# Auto-detection of system-first template requirement
# ======================================================================
Q35_TEMPLATE = (
"{%- for message in messages %}"
"{%- if message.role == 'system' %}"
"{%- if not loop.first %}"
"{{- raise_exception('System message must be at the beginning.') }}"
"{%- endif %}"
"{%- endif %}"
"{%- endfor %}"
)
class TestDetectMergeInlineSystem:
"""Verify _detect_merge_inline_system auto-detection.
Tests three scenarios:
1. Template with system-first guard (e.g. Qwen) → merge needed
2. Template without restrictions → no merge, cache-friendly
3. No template provided → safe default: merge
"""
def test_qwen_template_requires_merge(self):
"""Template with loop.first guard rejects mid-conversation system."""
assert (
AnthropicServingMessages._detect_merge_inline_system(Q35_TEMPLATE) is True
)
def test_no_restriction_no_merge(self):
"""Template without restriction accepts mid-conversation system."""
assert (
AnthropicServingMessages._detect_merge_inline_system(
"{%- for message in messages %}"
"{{- message.role }}: {{ message.content }}\n"
"{%- endfor %}"
)
is False
)
def test_no_template_defaults_merge(self):
"""No chat_template → conservative default: merge."""
assert AnthropicServingMessages._detect_merge_inline_system(None) is True
@@ -364,7 +364,7 @@ class MockVLLMServer:
await self._serve_task
def launch_mock_vllm(child_args: argparse.Namespace, env_updates: dict[str, str]):
def launch_mock_vllm(child_args: argparse.Namespace):
logger.info("Launching mock vLLM on port %s", child_args.port)
mock_vllm = MockVLLMServer(
port=child_args.port,
@@ -375,7 +375,7 @@ def launch_mock_vllm(child_args: argparse.Namespace, env_updates: dict[str, str]
def launch_mock_vllm_with_drain(
child_args: argparse.Namespace, env_updates: dict[str, str]
child_args: argparse.Namespace,
):
logger.info("Launching mock vLLM with 15s drain on port %s", child_args.port)
mock_vllm = MockVLLMServer(
+57 -44
View File
@@ -6,15 +6,22 @@ from typing import Final
import pytest
import schemathesis
from hypothesis import HealthCheck, settings
from schemathesis import GenerationConfig
from schemathesis.models import Case
from schemathesis import GenerationMode
from schemathesis.config import (
ChecksConfig,
CoveragePhaseConfig,
GenerationConfig,
PhasesConfig,
PositiveDataAcceptanceConfig,
ProjectConfig,
ProjectsConfig,
SchemathesisConfig,
)
from vllm.platforms import current_platform
from ...utils import RemoteOpenAIServer
schemathesis.experimental.OPEN_API_3_1.enable()
MODEL_NAME = "HuggingFaceTB/SmolVLM-256M-Instruct"
MAXIMUM_IMAGES = 2
_ROCM_TIMEOUT_MULTIPLIER = 3 if current_platform.is_rocm() else 1
@@ -44,21 +51,38 @@ def server():
@pytest.fixture(scope="module")
def get_schema(server):
# avoid generating null (\x00) bytes in strings during test case generation
return schemathesis.openapi.from_uri(
return schemathesis.openapi.from_url(
f"{server.url_root}/openapi.json",
generation_config=GenerationConfig(allow_x00=False),
config=SchemathesisConfig(
projects=ProjectsConfig(
default=ProjectConfig(
generation=GenerationConfig(
allow_x00=False,
modes=[GenerationMode.POSITIVE],
),
checks=ChecksConfig(
positive_data_acceptance=PositiveDataAcceptanceConfig(
enabled=False,
),
),
phases=PhasesConfig(
coverage=CoveragePhaseConfig(enabled=False),
),
),
),
),
)
schema = schemathesis.from_pytest_fixture("get_schema")
schema = schemathesis.pytest.from_fixture("get_schema")
@schemathesis.hook
def before_generate_case(context: schemathesis.hooks.HookContext, strategy):
def before_generate_case(context: schemathesis.HookContext, strategy):
op = context.operation
assert op is not None
def no_invalid_types(case: schemathesis.models.Case):
def no_invalid_types(case: schemathesis.Case):
"""
Skips tool_calls with `"type": "custom"` which schemathesis incorrectly
generates instead of the valid `"type": "function"`.
@@ -68,39 +92,25 @@ def before_generate_case(context: schemathesis.hooks.HookContext, strategy):
-d '{"messages": [{"role": "assistant", "tool_calls": [{"custom": {"input": "", "name": ""}, "id": "", "type": "custom"}]}]}' \
http://localhost:8000/v1/chat/completions
""" # noqa: E501
if hasattr(case, "body") and isinstance(case.body, dict):
if (
"messages" in case.body
and isinstance(case.body["messages"], list)
and len(case.body["messages"]) > 0
):
for message in case.body["messages"]:
if not isinstance(message, dict):
continue
if (
hasattr(case, "body")
and isinstance(case.body, dict)
and "messages" in case.body
and isinstance(case.body["messages"], list)
and len(case.body["messages"]) > 0
):
for message in case.body["messages"]:
if not isinstance(message, dict):
continue
tool_calls = message.get("tool_calls", [])
if isinstance(tool_calls, list):
for tool_call in tool_calls:
if isinstance(tool_call, dict):
if tool_call.get("type") != "function":
return False
if "custom" in tool_call:
return False
# Sometimes structured_outputs.grammar is generated to be empty
# Causing a server error in EBNF grammar parsing
# https://github.com/vllm-project/vllm/pull/22587#issuecomment-3195253421
structured_outputs = case.body.get("structured_outputs", {})
grammar = (
structured_outputs.get("grammar")
if isinstance(structured_outputs, dict)
else None
)
if grammar == "":
# Allow None (will be handled as no grammar)
# But skip empty strings
return False
tool_calls = message.get("tool_calls", [])
if isinstance(tool_calls, list):
for tool_call in tool_calls:
if isinstance(tool_call, dict):
if tool_call.get("type") != "function":
return False
if "custom" in tool_call:
return False
return True
@@ -108,7 +118,6 @@ def before_generate_case(context: schemathesis.hooks.HookContext, strategy):
@schema.parametrize()
@schema.override(headers={"Content-Type": "application/json"})
@settings(
deadline=LONG_TIMEOUT_SECONDS * 1000,
max_examples=50,
@@ -122,7 +131,7 @@ def before_generate_case(context: schemathesis.hooks.HookContext, strategy):
# generating large-but-valid request bodies before vLLM is called.
suppress_health_check=[HealthCheck.filter_too_much, HealthCheck.data_too_large],
)
def test_openapi_stateless(case: Case):
def test_openapi_stateless(case: schemathesis.Case):
key = (
case.operation.method.upper(),
case.operation.path,
@@ -151,4 +160,8 @@ def test_openapi_stateless(case: Case):
}.get(key, DEFAULT_TIMEOUT_SECONDS)
# No need to verify SSL certificate for localhost
case.call_and_validate(verify=False, timeout=timeout)
case.call_and_validate(
verify=False,
timeout=timeout,
headers={"Content-Type": "application/json"},
)
@@ -0,0 +1,65 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests that validation_exception_handler populates the `param` field
in its error response using the Pydantic error's `loc`, even when no
custom VLLMValidationError context is present.
Previously, `param` was only populated for errors carrying a custom
VLLMValidationError in their Pydantic `ctx`. Plain validation failures
(missing fields, wrong types) left `param` as None, even though the
field name was readily available from `error['loc']`.
"""
import json
from types import SimpleNamespace
import pytest
from fastapi.exceptions import RequestValidationError
from vllm.entrypoints.serve.utils.server_utils import validation_exception_handler
def _fake_request(log_error_stack: bool = False) -> SimpleNamespace:
"""Minimal stand-in for a FastAPI Request - just enough for the
handler to read req.app.state.args.log_error_stack."""
return SimpleNamespace(
app=SimpleNamespace(
state=SimpleNamespace(args=SimpleNamespace(log_error_stack=log_error_stack))
),
state=SimpleNamespace(), # no request_metadata -> hasattr(...) is False
)
class TestValidationErrorParamFallback:
"""Ensure `param` falls back to the Pydantic error's `loc` when no
custom VLLMValidationError context is present."""
@pytest.mark.parametrize(
("error_type", "msg"),
[
("missing", "Field required"),
("list_type", "Input should be a valid list"),
],
ids=["missing-field", "wrong-type"],
)
@pytest.mark.asyncio
async def test_param_falls_back_to_loc(self, error_type: str, msg: str):
errors = [{"type": error_type, "loc": ("body", "messages"), "msg": msg}]
exc = RequestValidationError(errors)
response = await validation_exception_handler(_fake_request(), exc)
body = json.loads(response.body)
assert body["error"]["param"] == "body.messages"
@pytest.mark.asyncio
async def test_param_fallback_does_not_crash_on_non_dict_error(self):
"""Schemathesis fuzzing found that errors[0] isn't always a dict.
The fallback must not crash in that case - it should just leave
param as None instead of raising."""
exc = RequestValidationError(["some unexpected non-dict error"])
response = await validation_exception_handler(_fake_request(), exc)
body = json.loads(response.body)
assert body["error"]["param"] is None
@@ -78,7 +78,16 @@ def test_gsm8k_correctness(config_filename):
"Skipping DeepSeek-V3.2 and DeepSeek-R1 on ROCm platforms "
"due to agent pool disk space issues and pod evictions."
)
if current_platform.is_rocm() and (
"Qwen3.5-35B-A3B-MXFP4-AITER-TP2" in config_filename.name
):
from vllm.platforms.rocm import on_gfx950
if not on_gfx950():
pytest.skip(
"Skipping Qwen3.5-35B-A3B-MXFP4-AITER-TP2 on non-GFX950 platforms. "
"The quantization scheme is not supported on non-GFX950 platforms."
)
# Parse server arguments from config (use shlex to handle quoted strings)
server_args_str = eval_config.get("server_args", "")
server_args = shlex.split(server_args_str) if server_args_str else []
+325 -1
View File
@@ -134,6 +134,7 @@ def _reference_index_topk(
topk: int,
init_blocks: int,
local_blocks: int,
sm_scale: float = 1.0,
) -> torch.Tensor:
total_q, num_idx_heads, _ = idx_q.shape
out = torch.full(
@@ -149,7 +150,7 @@ def _reference_index_topk(
num_blocks = (seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
pages = block_table[req_id, :num_blocks]
k = index_kv_cache[pages].reshape(num_blocks * BLOCK_SIZE, -1)
score = torch.einsum("qhd,kd->hqk", q.float(), k.float())
score = torch.einsum("qhd,kd->hqk", q.float(), k.float()) * sm_scale
q_pos = prefix_len + torch.arange(q_len, device=idx_q.device)
k_pos = torch.arange(k.shape[0], device=idx_q.device)
@@ -244,6 +245,270 @@ def test_prefill_index_topk_correctness():
_assert_topk_indices_equal_unordered(actual, expected)
# MSA indexer (SM100): fmha_sm100 OnlyScore for the per-block scores, then the
# Triton minimax_m3_index_topk for selection (no sparse_topk_select). Uses a
# deterministic construction (idx_q == 1, distinct e4m3-exact per-block values)
# so scores are strictly monotonic in the block id -> exact top-k agreement.
def _fmha_indexer_topk(
idx_q: torch.Tensor, # [total_q, H, 128] bf16/e4m3
index_cache: torch.Tensor, # [num_pages, 128, 128] bf16/e4m3
block_table: torch.Tensor,
q_lens: torch.Tensor,
seq_lens: torch.Tensor,
prefix_lens: torch.Tensor,
sm_scale: float,
topk: int,
) -> torch.Tensor:
"""Replicate MiniMaxM3IndexerMSAImpl's score path (single decode/prefill side)."""
from vllm.third_party.fmha_sm100.api import _fmha_sm100, _fmha_sm100_plan
num_idx_heads, head_dim = idx_q.shape[1], idx_q.shape[2]
nvp = [(s + 127) // 128 for s in seq_lens.tolist()]
kv_indices = torch.cat([block_table[r, : nvp[r]] for r in range(len(nvp))]).to(
torch.int32
)
qo = q_lens.cpu().to(torch.int32)
kv = seq_lens.cpu().to(torch.int32)
plan = _fmha_sm100_plan(
qo,
kv,
num_idx_heads,
num_kv_heads=1,
qo_offset=kv - qo,
page_size=128,
output_maxscore=True,
causal=True,
num_kv_splits=1,
)
k_pages = index_cache.view(index_cache.shape[0], 1, 128, head_dim)
_, max_score = _fmha_sm100(
idx_q,
k_pages,
k_pages,
plan,
kv_indices=kv_indices,
output_o=False,
output_maxscore=True,
sm_scale=sm_scale,
)
batch = q_lens.numel()
cu = torch.zeros(batch + 1, dtype=torch.int32, device=idx_q.device)
cu[1:] = q_lens.to(torch.int32).cumsum(0)
# max_score [H, k_tiles, total_q] -> transpose to [H, total_q, k_tiles].
return minimax_m3_index_topk(
max_score.transpose(1, 2),
cu,
prefix_lens.to(torch.int32),
int(q_lens.max()),
topk,
0, # init_blocks
0, # local_blocks
)
# e4m3-exact, strictly-increasing per-block values: with idx_q == 1 (also exact)
# the per-block scores are exact and distinct in BOTH bf16 and e4m3, so the fp8
# score path selects the same top-k as the reference (no quantization ties).
_E4M3_EXACT_VALUES = [
*range(1, 17), # 1..16 (step 1)
*range(18, 33, 2), # 18..32 (step 2)
*range(36, 65, 4), # 36..64 (step 4)
*range(72, 129, 8), # 72..128 (step 8)
]
@pytest.mark.skipif(
not current_platform.is_device_capability_family(100),
reason="fmha_sm100 indexer requires SM100 (Blackwell).",
)
@pytest.mark.parametrize("index_dtype", [torch.bfloat16, torch.float8_e4m3fn])
@pytest.mark.parametrize(
("q_lens", "prefix_lens"),
[
((4, 3), (2048, 2560)), # prefill: every token sees >= 16 causal blocks
((1, 1, 1), (2048, 3000, 4096)), # decode: one query token per request
],
)
def test_fmha_sm100_indexer_matches_reference(q_lens, prefix_lens, index_dtype):
torch.manual_seed(0)
num_idx_heads, head_dim = 4, HEAD_DIM
device = "cuda"
q_lens_t = torch.tensor(q_lens, device=device, dtype=torch.int32)
prefix_lens_t = torch.tensor(prefix_lens, device=device, dtype=torch.int32)
seq_lens = prefix_lens_t + q_lens_t
batch = len(q_lens)
max_blocks = (int(seq_lens.max()) + BLOCK_SIZE - 1) // BLOCK_SIZE
assert max_blocks <= len(_E4M3_EXACT_VALUES)
num_pages = batch * max_blocks
block_table = torch.randperm(num_pages, device=device, dtype=torch.int32).reshape(
batch, max_blocks
)
idx_q = torch.ones(
int(q_lens_t.sum()), num_idx_heads, head_dim, device=device, dtype=index_dtype
)
index_cache = torch.empty(
num_pages, BLOCK_SIZE, head_dim, device=device, dtype=index_dtype
)
for r in range(batch):
for b in range(max_blocks):
index_cache[block_table[r, b]] = float(_E4M3_EXACT_VALUES[b])
sm_scale = head_dim**-0.5
actual = _fmha_indexer_topk(
idx_q,
index_cache,
block_table,
q_lens_t,
seq_lens,
prefix_lens_t,
sm_scale,
TOPK,
)
expected = _reference_index_topk(
idx_q,
index_cache,
block_table,
q_lens_t,
seq_lens,
prefix_lens_t,
TOPK,
init_blocks=0,
local_blocks=0,
sm_scale=sm_scale,
)
_assert_topk_indices_equal_unordered(actual, expected)
# Full impl-level parity: drive both MiniMaxM3IndexerMSAImpl (fmha_sm100 score +
# Triton top-k) and MiniMaxM3IndexerTritonImpl through their real metadata
# builders on the SAME CommonAttentionMetadata + index cache, and assert the
# selected blocks agree. This exercises all the metadata the impl/kernels consume
# (decode/prefill split, cu_seqlens_q rebasing, prefix_lens, kv_indices gather,
# decode_pages split) -- a metadata bug on either side shifts the causal window
# or the block->page mapping and breaks the comparison.
@pytest.mark.skipif(
not current_platform.is_device_capability_family(100),
reason="fmha_sm100 indexer requires SM100 (Blackwell).",
)
@pytest.mark.parametrize("topk", [8, 16])
def test_msa_indexer_impl_matches_triton(topk, monkeypatch):
import vllm.models.minimax_m3.common.indexer as indexer_mod
from tests.v1.attention.utils import (
BatchSpec,
create_common_attn_metadata,
create_vllm_config,
)
from vllm.config import set_current_vllm_config
from vllm.forward_context import set_forward_context
from vllm.models.minimax_m3.common.indexer import (
MiniMaxM3IndexerTritonImpl,
MiniMaxM3IndexerTritonMetadataBuilder,
)
from vllm.models.minimax_m3.nvidia.indexer_msa import (
MiniMaxM3IndexerMSAImpl,
MiniMaxM3IndexerMSAMetadataBuilder,
)
torch.manual_seed(0)
device = torch.device("cuda")
num_idx_heads, head_dim = 4, HEAD_DIM
# TP=1: avoid requiring an initialized distributed group in a unit test.
monkeypatch.setattr(indexer_mod, "get_tensor_model_parallel_world_size", lambda: 1)
vllm_config = create_vllm_config(
block_size=BLOCK_SIZE, max_model_len=8192, max_num_batched_tokens=8192
)
vllm_config.model_config.hf_config.sparse_attention_config = {
"sparse_num_index_heads": num_idx_heads
}
# Decode-first mixed batch: 2 decode reqs (q_len 1) then 2 prefill reqs. Long
# prefixes so every token sees > TOPK causal blocks (non-trivial selection).
batch = BatchSpec(seq_lens=[2305, 2561, 2624, 2720], query_lens=[1, 1, 64, 96])
common = create_common_attn_metadata(
batch, BLOCK_SIZE, device, arange_block_indices=True
)
num_tokens = batch.compute_num_tokens()
# Deterministic index cache: distinct, monotonic per-logical-block values so
# the top-k is unambiguous (both kernels pick the same blocks, no fp ties).
block_table = common.block_table_tensor
num_pages = int(block_table.max().item()) + 1
index_cache = torch.zeros(
num_pages, BLOCK_SIZE, head_dim, device=device, dtype=DTYPE
)
for r, seq_len in enumerate(batch.seq_lens):
for b in range((seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE):
index_cache[block_table[r, b]] = float(b + 1)
index_q = torch.ones(
num_tokens, num_idx_heads * head_dim, device=device, dtype=DTYPE
)
spec = MLAAttentionSpec(
block_size=BLOCK_SIZE, num_kv_heads=1, head_size=head_dim, dtype=DTYPE
)
impl_kwargs = dict(
num_kv_heads=num_idx_heads,
scale=head_dim**-0.5,
topk_blocks=topk,
sparse_block_size=BLOCK_SIZE,
num_index_heads=num_idx_heads,
index_head_dim=head_dim,
init_blocks=0,
local_blocks=0,
)
with set_current_vllm_config(vllm_config):
msa_impl = MiniMaxM3IndexerMSAImpl(prefix="idx_msa", **impl_kwargs)
triton_impl = MiniMaxM3IndexerTritonImpl(prefix="idx_triton", **impl_kwargs)
msa_builder = MiniMaxM3IndexerMSAMetadataBuilder(
spec, [msa_impl.index_cache.prefix], vllm_config, device
)
triton_builder = MiniMaxM3IndexerTritonMetadataBuilder(
spec, [triton_impl.index_cache.prefix], vllm_config, device
)
# Both impls score against the same index keys.
msa_impl.index_cache.kv_cache = index_cache
triton_impl.index_cache.kv_cache = index_cache
# Exercise the shared persistent top-k buffer for BOTH impls: each must write
# decode ([:, :nd]) and prefill ([:, nd:]) into its buffer and return views.
# Separate buffers so the two forwards don't clobber each other.
nd = sum(q for q in batch.query_lens if q <= 1)
msa_impl.topk_indices_buffer = torch.full(
(num_idx_heads, num_tokens, topk), -2, dtype=torch.int32, device=device
)
triton_impl.topk_indices_buffer = torch.full(
(num_idx_heads, num_tokens, topk), -2, dtype=torch.int32, device=device
)
attn_metadata = {
msa_impl.index_cache.prefix: msa_builder.build(0, common),
triton_impl.index_cache.prefix: triton_builder.build(0, common),
}
with set_forward_context(attn_metadata, vllm_config):
msa_decode, msa_prefill = msa_impl(index_q)
tri_decode, tri_prefill = triton_impl(index_q)
assert msa_decode is not None and tri_decode is not None
assert msa_prefill is not None and tri_prefill is not None
_assert_topk_indices_equal_unordered(msa_decode, tri_decode)
_assert_topk_indices_equal_unordered(msa_prefill, tri_prefill)
# decode/prefill outputs are views into each impl's persistent buffer.
for impl, dec, pre in (
(msa_impl, msa_decode, msa_prefill),
(triton_impl, tri_decode, tri_prefill),
):
buf = impl.topk_indices_buffer
assert dec.data_ptr() == buf[:, :nd, :].data_ptr()
assert pre.data_ptr() == buf[:, nd:, :].data_ptr()
@pytest.mark.parametrize(
("decode_query_len", "max_decode_query_len"),
[
@@ -317,6 +582,65 @@ def test_decode_index_topk_correctness(
_assert_topk_indices_equal_unordered(actual, expected)
@pytest.mark.skipif(
not current_platform.is_device_capability_family(100),
reason="fp8 e4m3 indexer cache is the SM100 (MSA) path.",
)
@pytest.mark.parametrize("num_idx_heads", [1, 4])
def test_decode_index_topk_fp8(num_idx_heads: int):
"""The fp8 (e4m3) indexer cache feeds the Triton decode kernel on the MSA
path. The kernel must score in fp32 (no scaling) so its top-k matches a
reference computed from the dequantized fp8 values."""
torch.manual_seed(0)
topk, init_blocks, local_blocks, head_dim = 8, 0, 1, 128
decode_query_len = 1
active_seq_lens = torch.tensor((129, 1025, 4097), device="cuda", dtype=torch.int32)
q_lens = torch.full_like(active_seq_lens, decode_query_len)
prefix_lens = active_seq_lens - decode_query_len
batch = active_seq_lens.numel()
max_seq_len = int(active_seq_lens.max())
max_blocks = (max_seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
num_pages = batch * max_blocks
block_table = torch.randperm(num_pages, device="cuda", dtype=torch.int32).reshape(
batch, max_blocks
)
idx_q = torch.randn(
batch * decode_query_len, num_idx_heads, head_dim, device="cuda"
).to(torch.float8_e4m3fn)
index_kv_cache = torch.randn(num_pages, BLOCK_SIZE, head_dim, device="cuda").to(
torch.float8_e4m3fn
)
actual = minimax_m3_index_decode(
idx_q,
index_kv_cache,
block_table,
active_seq_lens,
max_seq_len=max_seq_len,
topk=topk,
init_blocks=init_blocks,
local_blocks=local_blocks,
num_kv_heads=num_idx_heads,
sm_scale=head_dim**-0.5,
decode_query_len=decode_query_len,
)
# Reference from the DEQUANTIZED fp8 values (the kernel computes the fp8 QK
# in fp32, so it must match an fp32 matmul of the same e4m3 values).
expected = _reference_index_topk(
idx_q.float(),
index_kv_cache.float(),
block_table,
q_lens,
active_seq_lens,
prefix_lens,
topk,
init_blocks,
local_blocks,
head_dim**-0.5,
)
_assert_topk_indices_equal_unordered(actual, expected)
# Sparse attention kernels.
def _reference_sparse_attn(
q: torch.Tensor,
+57 -2
View File
@@ -5,8 +5,9 @@ Tests for the FlashInfer TRTLLM NvFP4 MoE backend
(`TrtLlmNvFp4ExpertsModular`).
Covers the activations the wrapper claims to support — SiLU, RELU^2 (non-gated),
and GELU — including a Gemma4-shaped case (128 experts, top-k 8,
intermediate_size 704) that exercises the non-256-aligned padding path.
GELU, and clamped SwiGLU-OAI (MiniMax-M3) — including a Gemma4-shaped case
(128 experts, top-k 8, intermediate_size 704) that exercises the non-256-aligned
padding path.
"""
import pytest
@@ -80,6 +81,29 @@ if _CLAMP_OP_NAME not in op_registry:
SILU_WITH_CLAMP = op_registry[_CLAMP_OP_NAME]
# Clamped SwiGLU-OAI (MiniMax-M3): non-default alpha/beta so the kernel must
# honor gemm1_alpha (raw) and gemm1_beta (folded by g1_alphas), not just clamp.
_SWIGLU_ALPHA = 1.702
_SWIGLU_BETA = 1.0
_OAI_OP_NAME = "test_swigluoai_with_clamp"
if _OAI_OP_NAME not in op_registry:
@CustomOp.register(_OAI_OP_NAME)
class _SwigluOAIWithClampTest(SiluAndMulWithClamp):
custom_op_name = _OAI_OP_NAME
def __init__(self, *, compile_native: bool = True) -> None:
super().__init__(
_SWIGLU_LIMIT,
alpha=_SWIGLU_ALPHA,
beta=_SWIGLU_BETA,
compile_native=compile_native,
)
SWIGLUOAI_REF = op_registry[_OAI_OP_NAME]
ACTIVATION_CASES = [
pytest.param(MoEActivation.SILU, MoEActivation.SILU, None, id="silu"),
@@ -91,6 +115,12 @@ ACTIVATION_CASES = [
id="relu2_no_mul",
),
pytest.param(MoEActivation.GELU, MoEActivation.GELU, None, id="gelu"),
pytest.param(
MoEActivation.SWIGLUOAI_UNINTERLEAVE,
SWIGLUOAI_REF,
_SWIGLU_LIMIT,
id="swigluoai_uninterleave",
),
]
@@ -148,6 +178,10 @@ def test_trtllm_fp4_moe_no_graph(
is_scale_swizzled=False,
)
quant_config.gemm1_clamp_limit = swiglu_limit
is_oai = activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE
if is_oai:
quant_config.gemm1_alpha = _SWIGLU_ALPHA
quant_config.gemm1_beta = _SWIGLU_BETA
if swiglu_limit is not None:
assert quant_config.g1_alphas is not None
assert quant_config.a2_gscale is not None
@@ -192,6 +226,27 @@ def test_trtllm_fp4_moe_no_graph(
fake_layer.w2_input_scale = torch.ones_like(quant_config.g2_alphas)
trtllm_inner.process_weights_after_loading(fake_layer)
if is_oai:
# alpha stays raw; beta and clamp are folded by g1_alphas
# (== _LARGE_OUTPUT1_SCALE here), so the fold is load-bearing.
assert torch.allclose(
trtllm_inner.gemm1_alpha,
torch.full_like(trtllm_inner.gemm1_alpha, _SWIGLU_ALPHA),
)
assert torch.allclose(
trtllm_inner.gemm1_beta,
torch.full_like(
trtllm_inner.gemm1_beta, _SWIGLU_BETA / _LARGE_OUTPUT1_SCALE
),
)
assert torch.allclose(
trtllm_inner.gemm1_clamp_limit,
torch.full_like(
trtllm_inner.gemm1_clamp_limit,
_SWIGLU_LIMIT / _LARGE_OUTPUT1_SCALE,
),
)
trtllm_experts = mk.FusedMoEKernel(
maybe_make_prepare_finalize(
moe=moe_config,
@@ -345,6 +345,63 @@ def test_per_token_group_quant_fp8_packed_zero_fills_padded_output_q(
)
@pytest.mark.skipif(
not current_platform.is_cuda_alike(),
reason="packed FP8 per-token-group quant kernel requires a CUDA-alike GPU",
)
def test_per_token_group_quant_fp8_packed_large_mn():
"""Regression test for https://github.com/vllm-project/vllm/issues/45099.
Some background: gridDim.x and gridDim.y have different limits of 2^31 - 1 and
2^16 - 1, respectively.
Prior code introduced a bug where it incorrectly assumed grid.x and y both have
2^31 - 1 limits and mixed them up, which doesn't surface until the kernel is
launched with a large mn that exceeds grid.y limit (2^16 - 1).
This issue doesn't surface often because each forward pass only processes a
bounded token batch, not the full context.
Quantizing tensors with more rows than that will fail at launch with
"CUDA error: invalid argument".
This is a differential test that compares fp8 output against Triton output
reference when token size sits just above the gridDim.y 2^16 - 1 limit.
"""
device = "cuda"
group_size = 128
# hidden 2048 -> 2048/128 = 16 groups per row -> kx=16, ry=1: one grid row per mn
# row, so any mn > 65535 overflowed grid.y before the fix.
num_tokens, hidden_dim = 65537, 2048
torch.manual_seed(42)
x = torch.randn((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16) * 8
out_q, out_s_packed = fp8_utils.per_token_group_quant_fp8_packed_for_deepgemm(
x,
group_size=group_size,
use_ue8m0=True,
)
with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False):
ref_q, ref_s = fp8_utils.per_token_group_quant_fp8(
x, group_size, use_ue8m0=True
)
assert torch.equal(out_q, ref_q), "Quantized output mismatch"
# Vectorized packed-scale check; the per-element loop used by the smaller
# tests is too slow at this size. groups_per_row is a multiple of 4 here,
# so there is no K padding and the packed view lines up.
mn = num_tokens
groups_per_row = hidden_dim // group_size
k_num_packed = (groups_per_row + 3) // 4
assert groups_per_row % 4 == 0
ref_exponents = (ref_s.reshape(mn, groups_per_row).view(torch.int32) >> 23) & 0xFF
exp = ref_exponents.view(mn, k_num_packed, 4)
expected = (
exp[..., 0] | (exp[..., 1] << 8) | (exp[..., 2] << 16) | (exp[..., 3] << 24)
)
assert torch.equal(out_s_packed.cpu(), expected.cpu()), "Packed scale mismatch"
@pytest.mark.parametrize("shape", [(32, 128), (64, 256), (16, 512)])
@pytest.mark.parametrize("group_size", [64, 128])
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
@@ -60,8 +60,10 @@ def test_rocm_compressed_tensors_w8a8(
vllm_runner, example_prompts, model_path, max_tokens, num_logprobs
):
dtype = "bfloat16"
with vllm_runner(model_path, dtype=dtype) as vllm_model:
# Pin to TRITON_ATTN, see https://github.com/vllm-project/vllm/issues/46179
with vllm_runner(
model_path, dtype=dtype, attention_backend="TRITON_ATTN"
) as vllm_model:
vllm_model.generate_greedy_logprobs(example_prompts, max_tokens, num_logprobs)
@@ -278,3 +278,99 @@ def test_sparse_full(num_tokens, block_size, kv_cache_dtype):
torch.testing.assert_close(
index_cache.view(-1, HEAD_DIM), expected_index_cache, rtol=0, atol=0
)
# ── Test 3: fp8 (e4m3) index outputs ─────────────────────────────────────────
# The fp8 score path stores index_q and the index-K cache as e4m3 while q/k/v +
# q_out stay bf16. Asserts: (1) q/k/v/q_out are bit-identical to the bf16 run
# (the index dtype must not perturb the main branch), and (2) the e4m3 index
# outputs dequantize close to the bf16 reference.
@pytest.mark.skipif(
not torch.cuda.is_available() or torch.cuda.get_device_capability() < (8, 9),
reason="e4m3 conversion requires CUDA SM89+.",
)
@pytest.mark.parametrize("num_tokens", [1, 7, 64, 513])
@pytest.mark.parametrize("block_size", [16, 64])
def test_sparse_full_fp8_index(num_tokens, block_size):
torch.manual_seed(1)
device, dtype, eps = "cuda", torch.bfloat16, 1e-6
base, max_pos = 5_000_000.0, 4096
num_heads, num_kv_heads, num_idx_heads = 16, 4, 4
q_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
k_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
iq_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
ik_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
cos_sin = make_cos_sin_cache(max_pos, ROTARY_DIM, base, dtype, device)
positions = torch.randint(
0, max_pos, (num_tokens,), dtype=torch.int64, device=device
)
qsz, kvsz = num_heads * HEAD_DIM, num_kv_heads * HEAD_DIM
iqsz, iksz = num_idx_heads * HEAD_DIM, HEAD_DIM
qkv0 = torch.randn(
num_tokens, qsz + 2 * kvsz + iqsz + iksz, dtype=dtype, device=device
)
num_blocks = (num_tokens + block_size - 1) // block_size + 1
slot_mapping = torch.randperm(
num_blocks * block_size, dtype=torch.int64, device=device
)[:num_tokens]
index_slot_mapping = torch.roll(slot_mapping, shifts=1)
def run(index_dtype):
qkv = qkv0.clone()
kv_cache = torch.zeros(
num_blocks,
2,
block_size,
num_kv_heads,
HEAD_DIM,
dtype=dtype,
device=device,
)
index_cache = torch.zeros(
num_blocks, block_size, HEAD_DIM, dtype=index_dtype, device=device
)
q_out = torch.empty(num_tokens, qsz, dtype=dtype, device=device)
index_q = torch.empty(num_tokens, iqsz, dtype=index_dtype, device=device)
ops.fused_minimax_m3_qknorm_rope_kv_insert(
qkv,
q_w,
k_w,
cos_sin,
positions,
num_heads,
num_kv_heads,
ROTARY_DIM,
eps,
iq_w,
ik_w,
num_idx_heads,
slot_mapping,
index_slot_mapping,
kv_cache,
index_cache,
block_size,
q_out,
index_q,
)
return qkv, kv_cache, index_cache, q_out, index_q
qkv_bf, kvc_bf, idxc_bf, qo_bf, iq_bf = run(torch.bfloat16)
qkv_fp, kvc_fp, idxc_fp, qo_fp, iq_fp = run(torch.float8_e4m3fn)
assert iq_fp.dtype == torch.float8_e4m3fn
assert idxc_fp.dtype == torch.float8_e4m3fn
# (1) The main branch (q/k/v in qkv, q_out, kv cache) must be bit-identical:
# the index output dtype must not perturb anything else.
torch.testing.assert_close(qo_fp, qo_bf, rtol=0, atol=0)
torch.testing.assert_close(qkv_fp, qkv_bf, rtol=0, atol=0)
torch.testing.assert_close(kvc_fp, kvc_bf, rtol=0, atol=0)
# (2) Dequantized e4m3 index outputs match the bf16 reference within fp8 ulp.
torch.testing.assert_close(iq_fp.float(), iq_bf.float(), rtol=0.13, atol=0.05)
torch.testing.assert_close(idxc_fp.float(), idxc_bf.float(), rtol=0.13, atol=0.05)
+124
View File
@@ -482,3 +482,127 @@ def test_kernels_hidden_size(
seq_length=128,
add_inputs=True,
)
@pytest.mark.parametrize("device", DEVICES)
def test_add_lora_fused_moe_early_exit(device):
"""
Ensures add_lora_fused_moe does not invoke the LoRA kernel or
modify the output tensor when no_lora_flag_cpu is True
"""
from types import SimpleNamespace
from vllm.lora.punica_wrapper.punica_gpu import PunicaWrapperGPU
torch.set_default_device(device)
torch.accelerator.set_device_index(device)
max_loras, num_tokens = 4, 16
num_experts, top_k, max_lora_rank = 8, 2, 16
K, N = 256, 128
# build PunicaWrapperGPU with minimal lora_config mock
lora_config = SimpleNamespace(
max_loras=max_loras,
specialize_active_lora=False,
)
wrapper = PunicaWrapperGPU(
max_num_batched_tokens=num_tokens,
max_batches=num_tokens,
device=device,
lora_config=lora_config,
)
# simulate a prior LoRA batch so the internal mapping is
# populated with stale LoRA IDs
lora_mapping = torch.zeros(
num_tokens,
dtype=torch.int32,
device=device,
)
lora_mapping[:8] = 1
lora_mapping[8:] = 2
wrapper.token_mapping_meta.prepare_tensors(lora_mapping)
# simulate a base-model batch (all -1)
base_mapping = torch.full(
(num_tokens,),
-1,
dtype=torch.int32,
device=device,
)
wrapper.token_mapping_meta.prepare_tensors(base_mapping)
assert wrapper.token_mapping_meta.no_lora_flag_cpu[0].item() is True
# dummy tensors for add_lora_fused_moe
y = torch.rand(num_tokens, top_k, N, dtype=torch.bfloat16, device=device)
y_snapshot = y.clone()
x = torch.rand(num_tokens, K, dtype=torch.bfloat16, device=device)
lora_a_stacked = (
torch.rand(
max_loras,
num_experts,
max_lora_rank,
K,
dtype=torch.bfloat16,
device=device,
),
)
lora_b_stacked = (
torch.rand(
max_loras,
num_experts,
N,
max_lora_rank,
dtype=torch.bfloat16,
device=device,
),
)
topk_weights = torch.ones(
num_tokens,
top_k,
dtype=torch.float32,
device=device,
)
adapter_enabled = torch.ones(
max_loras + 1,
dtype=torch.int32,
device=device,
)
shrink_config = expand_config = {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 32,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"NUM_WARPS": 4,
"NUM_STAGES": 3,
"SPLIT_K": 1,
}
# call add_lora_fused_moe - the early exit should prevent any
# modification to the output
wrapper.add_lora_fused_moe(
y=y,
x=x,
lora_a_stacked=lora_a_stacked,
lora_b_stacked=lora_b_stacked,
topk_weights=topk_weights,
sorted_token_ids=None,
expert_ids=torch.zeros(
num_tokens * top_k,
dtype=torch.int32,
device=device,
),
num_tokens_post_padded=None,
max_lora_rank=max_lora_rank,
top_k_num=top_k,
shrink_config=shrink_config,
expand_config=expand_config,
adapter_enabled=adapter_enabled,
)
assert torch.equal(y, y_snapshot), (
"add_lora_fused_moe modified output tensor despite no_lora_flag_cpu=True"
)
@@ -130,8 +130,12 @@ def test_models(
monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
if model == "TitanML/tiny-mixtral":
# Untrained model: near-uniform logits make argmax sensitive to
# AITER's bfloat16 rounding error in plain rms_norm.
# AITER's bfloat16 rounding error. Route the plain rms_norm and the
# fused MoE (whose near-uniform router logits flip expert selection
# under ~1 ULP drift) through the native kernels for this model.
# See ROCm/aiter#3806 for the tracking issue and minimal repro.
monkeypatch.setenv("VLLM_ROCM_USE_AITER_RMSNORM", "0")
monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "0")
elif use_rocm_aiter and model not in AITER_MODEL_LIST:
# Skip model that are not using AITER tests.
# When more AITER kernels are added, this list will not be
@@ -25,7 +25,7 @@ TEST_IMAGE_NAMES = [
]
MAX_MODEL_LEN = 8192
REQUESTS_PER_ROUND = 4
WARMUP_ROUNDS = 1
WARMUP_ROUNDS = 2
MEASURED_ROUNDS = 16
GPU_GROWTH_THRESHOLD_MIB = 0
CPU_PEAK_GROWTH_THRESHOLD_MIB = 0
@@ -2,6 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
import torch.nn.functional as F
import transformers.utils
from PIL import Image
@@ -52,6 +53,7 @@ def _get_cherry_blossom_image() -> Image.Image:
)
@torch.inference_mode()
def _run_test(
hf_runner: type[HfRunner],
vllm_runner: type[VllmRunner],
@@ -92,3 +92,49 @@ def test_processor_num_frames_timestamp(
assert len(video_phs) == 1, (
f"Expected exactly 1 video placeholder, got {len(video_phs)}"
)
@pytest.mark.parametrize("model_id", [MODEL_ID])
@pytest.mark.parametrize("num_videos", [2, 4])
def test_processor_multi_video(
model_id: str,
num_videos: int,
) -> None:
"""Verify that multi-video processing produces correct placeholders.
This exercises the token-level replacement path in
``_call_hf_processor`` which avoids the quadratic text-level
prompt expansion.
"""
ctx = build_model_context(
model_id,
limit_mm_per_prompt={"image": 0, "video": num_videos},
)
processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
prompt = "<|vision_start|><|video_pad|><|vision_end|>" * num_videos
mm_data = {"video": [_build_video_mm_data(num_frames=8)["video"][0]] * num_videos}
processed = processor(
prompt,
mm_items=processor.info.parse_mm_data(mm_data),
hf_processor_mm_kwargs={"num_frames": 8},
)
token_ids = processed["prompt_token_ids"]
assert len(token_ids) > 0
video_phs = processed["mm_placeholders"].get("video", [])
assert len(video_phs) == num_videos, (
f"Expected {num_videos} video placeholders, got {len(video_phs)}"
)
# All placeholders should have the same length (same video params)
# and must not overlap.
lengths = {ph.length for ph in video_phs}
assert len(lengths) == 1, f"Placeholder lengths differ: {lengths}"
for i in range(1, len(video_phs)):
prev_end = video_phs[i - 1].offset + video_phs[i - 1].length
assert video_phs[i].offset >= prev_end, (
f"Placeholder {i} overlaps with placeholder {i - 1}"
)
+3
View File
@@ -96,6 +96,9 @@ class MockTokenizer:
return "".join(parts)
CHUNK_SIZES = [1, 2, 3, 5, 11, 23, None]
def make_mock_tokenizer(sample: Sample) -> MockTokenizer:
"""Build a mock tokenizer from a sample's vocab and token data."""
return MockTokenizer(
@@ -19,6 +19,7 @@ import pytest
from pydantic import TypeAdapter
from tests.parser.engine.replay_harness import (
CHUNK_SIZES,
MockTokenizer,
assert_parse_output,
collect_output,
@@ -113,8 +114,6 @@ _PAIRINGS = _discover_pairings()
_ALL_SAMPLES = [(p.parser_cls, s) for p in _PAIRINGS for s in p.samples]
CHUNK_SIZES = [1, 2, 3, 5, 11, 23, None]
@pytest.mark.parametrize("chunk_size", CHUNK_SIZES, ids=lambda c: f"chunk={c}")
@pytest.mark.parametrize(
@@ -0,0 +1,181 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Regression test for U+FFFD leak at reasoning→content transition.
When byte-fallback tokens span the reasoning/content boundary,
decoding isolated content-side token IDs via tokenizer.decode()
produces U+FFFD (Unicode replacement character). The fix flushes
the reasoning parser's engine lexer instead.
Reproduces the bug at various chunk sizes and validates that the
fix prevents U+FFFD from leaking into streamed content.
"""
from __future__ import annotations
import pytest
from tests.parser.engine.replay_harness import (
CHUNK_SIZES,
MockTokenizer,
collect_output,
replay_streaming,
)
from vllm.parser.abstract_parser import DelegatingParser
from vllm.parser.engine.registered_adapters import (
Glm47MoeParserReasoningAdapter,
Glm47MoeParserToolAdapter,
Qwen3ParserReasoningAdapter,
Qwen3ParserToolAdapter,
)
class ByteFallbackMockTokenizer(MockTokenizer):
"""MockTokenizer that returns U+FFFD for specified token IDs.
Simulates byte-fallback tokenizer behavior where isolated
partial-byte tokens decode to the Unicode replacement character.
"""
def __init__(
self,
vocab: dict[str, int],
tokens: list[tuple[int, str]],
ufffd_token_ids: set[int],
) -> None:
super().__init__(vocab, tokens)
self._ufffd_token_ids = frozenset(ufffd_token_ids)
def decode(self, ids: list[int], skip_special_tokens: bool = False) -> str:
parts: list[str] = []
for tid in ids:
if skip_special_tokens and tid in self._special_ids:
continue
if tid in self._ufffd_token_ids:
parts.append("")
else:
text = self._token_decode_map.get(tid, f"?{tid}?")
parts.append(text)
return "".join(parts)
# ── Model-specific DelegatingParser subclasses ───────────────────────
class _Glm47Delegating(DelegatingParser):
reasoning_parser_cls = Glm47MoeParserReasoningAdapter
tool_parser_cls = Glm47MoeParserToolAdapter
class _Qwen3Delegating(DelegatingParser):
reasoning_parser_cls = Qwen3ParserReasoningAdapter
tool_parser_cls = Qwen3ParserToolAdapter
# ── Shared test data ─────────────────────────────────────────────────
_SHARED_TOKENS: list[tuple[int, str]] = [
(100, "Let me"),
(101, " think"),
(102, " about"),
(103, " Samsung."),
(51, "</think>"),
(200, "삼성"),
(201, "전자의"),
(202, " 주가를"),
(203, " 분석합니다."),
]
_SHARED_UFFFD_IDS: set[int] = {200}
EXPECTED_REASONING = "Let me think about Samsung."
EXPECTED_CONTENT = "삼성전자의 주가를 분석합니다."
_MODEL_CONFIGS = [
pytest.param(
{
"<think>": 50,
"</think>": 51,
"<tool_call>": 60,
"</tool_call>": 61,
"<arg_key>": 62,
"</arg_key>": 63,
"<arg_value>": 64,
"</arg_value>": 65,
},
_Glm47Delegating,
id="glm47",
),
pytest.param(
{
"<think>": 50,
"</think>": 51,
"<tool_call>": 60,
"</tool_call>": 61,
},
_Qwen3Delegating,
id="qwen3",
),
]
# ── Tests ────────────────────────────────────────────────────────────
class TestUfffdReasoningTransition:
"""U+FFFD must not appear at the reasoning→content transition."""
@pytest.mark.parametrize("vocab,delegating_cls", _MODEL_CONFIGS)
@pytest.mark.parametrize("chunk_size", CHUNK_SIZES, ids=lambda c: f"chunk={c}")
def test_no_ufffd(self, chunk_size, vocab, delegating_cls):
tokenizer = ByteFallbackMockTokenizer(vocab, _SHARED_TOKENS, _SHARED_UFFFD_IDS)
parser = delegating_cls(tokenizer)
deltas = replay_streaming(
parser,
_SHARED_TOKENS,
chunk_size=chunk_size,
finished_on_last=True,
)
output = collect_output(deltas)
assert "" not in output.content, (
f"U+FFFD leaked into content: {output.content!r}"
)
assert output.content == EXPECTED_CONTENT
assert output.reasoning == EXPECTED_REASONING
def test_byte_fallback_tokenizer_produces_ufffd(self):
"""Validate the fixture: decode() returns U+FFFD for isolated
byte-fallback token IDs, proving the old code path would leak."""
vocab = dict(_MODEL_CONFIGS[0].values[0])
tokenizer = ByteFallbackMockTokenizer(vocab, _SHARED_TOKENS, _SHARED_UFFFD_IDS)
assert tokenizer.decode([200]) == ""
@pytest.mark.parametrize("chunk_size", CHUNK_SIZES, ids=lambda c: f"chunk={c}")
def test_multiple_ufffd_tokens_at_boundary(self, chunk_size):
"""Multiple consecutive byte-fallback tokens at the boundary."""
tokens: list[tuple[int, str]] = [
(100, "Reasoning."),
(51, "</think>"),
(200, ""),
(201, ""),
(202, "전자"),
]
ufffd_ids: set[int] = {200, 201}
vocab = dict(_MODEL_CONFIGS[0].values[0])
tokenizer = ByteFallbackMockTokenizer(vocab, tokens, ufffd_ids)
parser = _Glm47Delegating(tokenizer)
deltas = replay_streaming(
parser,
tokens,
chunk_size=chunk_size,
finished_on_last=True,
)
output = collect_output(deltas)
assert "" not in output.content, (
f"U+FFFD leaked into content: {output.content!r}"
)
assert output.content == "삼성전자"
assert output.reasoning == "Reasoning."
+17 -2
View File
@@ -143,6 +143,12 @@ SCENARIOS: list[Scenario] = [
tool_calls=[_READ_TOOL],
after_tool_response=True,
),
Scenario(
id="empty-tool-block",
description="Empty tool block followed by content (edge case recovery)",
content="Content after empty tools.",
tool_calls=[],
),
]
@@ -344,8 +350,11 @@ def _qwen3_segments(scenario: Scenario) -> list[tuple[str, bool]]:
segs: list[tuple[str, bool]] = []
if scenario.reasoning is not None:
segs.append((scenario.reasoning, False))
if scenario.content is not None or scenario.tool_calls:
if scenario.content is not None or scenario.tool_calls is not None:
segs.append(("</think>", True))
if scenario.tool_calls is not None and not scenario.tool_calls:
segs.append(("<tool_call>", True))
segs.append(("</tool_call>", True))
if scenario.content is not None:
segs.append((scenario.content, False))
if scenario.tool_calls:
@@ -437,8 +446,11 @@ def _minimax_m2_segments(scenario: Scenario) -> list[tuple[str, bool]]:
segs: list[tuple[str, bool]] = []
if scenario.reasoning is not None:
segs.append((scenario.reasoning, False))
if scenario.content is not None or scenario.tool_calls:
if scenario.content is not None or scenario.tool_calls is not None:
segs.append(("</think>", True))
if scenario.tool_calls is not None and not scenario.tool_calls:
segs.append(("<minimax:tool_call>", True))
segs.append(("</minimax:tool_call>", True))
if scenario.content is not None:
segs.append((scenario.content, False))
if scenario.tool_calls:
@@ -534,6 +546,9 @@ def _gemma4_segments(scenario: Scenario) -> list[tuple[str, bool]]:
segs.append((_GEMMA4_THOUGHT_PREFIX, False))
segs.append((scenario.reasoning, False))
segs.append(("<channel|>", True))
if scenario.tool_calls is not None and not scenario.tool_calls:
segs.append(("<|tool_call>", True))
segs.append(("<tool_call|>", True))
if scenario.content is not None:
segs.append((scenario.content, False))
if scenario.tool_calls:
@@ -18,6 +18,7 @@ from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import (
FP8ScaledMMLinearLayerConfig,
)
from vllm.model_executor.layers.quantization.utils.quant_utils import (
get_fp8_min_max,
kFp8DynamicTokenSym,
kFp8StaticChannelSym,
kFp8StaticTensorSym,
@@ -309,7 +310,7 @@ def test_hipb_mm_kernel_forward_accuracy(enable_hipb_mm_kernel):
_check_bpreshuffle_runtime_support(weight_shape, num_tokens=num_tokens)
fp8_dtype = current_platform.fp8_dtype()
fp8_max = torch.finfo(fp8_dtype).max
fp8_max = get_fp8_min_max()[1]
device = torch.device("cuda")
# Build a bf16 weight and quantize per output channel (one scale per row).
+135
View File
@@ -0,0 +1,135 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for contiguous KV cache packing in _get_kv_cache_config_deepseek_v4."""
from unittest.mock import MagicMock
import pytest
import torch
from vllm.v1.core.kv_cache_utils import _get_kv_cache_config_deepseek_v4
from vllm.v1.kv_cache_interface import (
KVCacheGroupSpec,
MLAAttentionSpec,
UniformTypeKVCacheSpecs,
)
def _make_mla_spec(page_size: int, block_size: int = 256) -> MLAAttentionSpec:
return MLAAttentionSpec(
block_size=block_size,
num_kv_heads=1,
head_size=512,
dtype=torch.uint8,
page_size_padded=page_size,
cache_dtype_str="fp8_ds_mla",
model_version="deepseek_v4",
alignment=576,
)
def _make_groups(n_c4, n_c128, n_swa):
PS_C4_MLA = 37440
PS_C4_IDX = 8640
PS_C128 = 1728
PS_SWA = 37440
mla_specs = {}
for i in range(n_c4):
mla_specs[f"c4_mla.{i}"] = _make_mla_spec(PS_C4_MLA)
mla_specs[f"c4_idx.{i}"] = _make_mla_spec(PS_C4_IDX)
for i in range(n_c128):
mla_specs[f"c128_mla.{i}"] = _make_mla_spec(PS_C128)
mla_group = KVCacheGroupSpec(
layer_names=list(mla_specs.keys()),
kv_cache_spec=UniformTypeKVCacheSpecs(block_size=256, kv_cache_specs=mla_specs),
)
swa_specs = {}
for i in range(n_swa):
swa_specs[f"swa.{i}"] = _make_mla_spec(PS_SWA)
swa_group = KVCacheGroupSpec(
layer_names=list(swa_specs.keys()),
kv_cache_spec=UniformTypeKVCacheSpecs(block_size=256, kv_cache_specs=swa_specs),
)
return [mla_group, swa_group]
def _mock_vllm_config():
config = MagicMock()
config.cache_config.num_gpu_blocks_override = None
return config
def _run(n_c4=3, n_c128=2, n_swa=5, mem=100 * 1024 * 1024):
groups = _make_groups(n_c4, n_c128, n_swa)
return _get_kv_cache_config_deepseek_v4(_mock_vllm_config(), groups, mem)
def _page_sizes_by_layer(
groups: list[KVCacheGroupSpec],
) -> dict[str, int]:
page_sizes = {}
for group in groups:
specs = group.kv_cache_spec.kv_cache_specs
for layer_name in group.layer_names:
page_sizes[layer_name] = specs[layer_name].page_size_bytes
return page_sizes
class TestInterleavedPacking:
def test_all_tensors_have_block_stride(self):
_, tensors = _run()
for t in tensors:
assert t.block_stride > 0
def test_all_tensors_share_same_size(self):
_, tensors = _run()
sizes = set(t.size for t in tensors)
assert len(sizes) == 1
assert sizes.pop() > 0
def test_offsets_within_one_block(self):
_, tensors = _run()
for t in tensors:
assert t.offset < t.block_stride
def test_all_layers_accounted_for(self):
n_c4, n_c128, n_swa = 5, 4, 7
_, tensors = _run(n_c4=n_c4, n_c128=n_c128, n_swa=n_swa)
all_names = set()
for t in tensors:
all_names.update(t.shared_by)
expected = n_c4 * 2 + n_c128 + n_swa
assert len(all_names) == expected
def test_strided_views_are_independent(self):
groups = _make_groups(n_c4=3, n_c128=2, n_swa=5)
page_sizes = _page_sizes_by_layer(groups)
num_blocks, tensors = _get_kv_cache_config_deepseek_v4(
_mock_vllm_config(), groups, 100 * 1024 * 1024
)
backing = torch.zeros(tensors[0].size, dtype=torch.uint8)
views = []
for t in tensors:
page_size = page_sizes[t.shared_by[0]]
v = torch.as_strided(
backing,
size=(num_blocks, page_size),
stride=(t.block_stride, 1),
storage_offset=t.offset,
)
views.append(v)
for i, v in enumerate(views):
v.fill_(i + 1)
for i, v in enumerate(views):
assert (v == i + 1).all(), f"View {i} was corrupted"
if __name__ == "__main__":
pytest.main([__file__, "-v"])
@@ -4,9 +4,16 @@
import pytest
from vllm import LLM, SamplingParams
from vllm.platforms import current_platform
from ....utils import create_new_process_for_each_test
if current_platform.is_rocm():
pytest.skip(
"Cascade attention backends FLASH_ATTN and FLASHINFER are notsupported on ROCm",
allow_module_level=True,
)
@create_new_process_for_each_test()
@pytest.mark.parametrize("attn_backend", ["FLASH_ATTN", "FLASHINFER"])
@@ -1,6 +1,8 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import threading
import time
from unittest.mock import MagicMock, patch
from vllm.config import set_current_vllm_config
@@ -406,7 +408,9 @@ def test_lookup_key_client_lookup_prepends_typed_tag():
fake_socket = mock_make_socket.return_value
fake_socket.recv.return_value = (5).to_bytes(4, "big")
assert client.lookup(token_len=128, block_hashes=[]) == 5
# Blocking lookup (non_block defaults to False) runs on the executor and
# returns the resolved hit length.
assert client.lookup("req0", token_len=128, block_hashes=[]) == 5
sent_frames = fake_socket.send_multipart.call_args[0][0]
assert sent_frames[0] == protocol.LOOKUP_MSG
@@ -435,6 +439,127 @@ def test_lookup_key_client_reset_uses_typed_protocol():
assert client.reset() is False
def _poll_lookup(client, req_id, token_len=128, block_hashes=(), timeout=5.0):
"""Drive non-blocking lookup until the executor completes it."""
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
result = client.lookup(req_id, token_len, list(block_hashes), non_block=True)
if result is not None:
return result
time.sleep(0.005)
return None
def _gated_recv(gate: threading.Event, value: int):
"""Mock recv side-effect that blocks until ``gate`` is set, so the
executor's lookup can be held pending deterministically."""
def recv():
gate.wait()
return value.to_bytes(4, "big")
return recv
def test_lookup_key_client_non_block_lookup_async():
"""Non-blocking lookup defers to the executor: None first, hit once the
Future resolves."""
vllm_config = _make_vllm_config()
with patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
"worker.make_zmq_socket"
) as mock_make_socket:
client = worker.LookupKeyClient(vllm_config)
fake_socket = mock_make_socket.return_value
# Hold the executor's lookup pending until we release the gate.
gate = threading.Event()
fake_socket.recv.side_effect = _gated_recv(gate, 7)
# First query submits the lookup and returns None while it is in flight.
assert client.lookup("req1", 128, [], non_block=True) is None
# Release the executor; a later poll returns the hit length.
gate.set()
assert _poll_lookup(client, "req1") == 7
# Future is consumed (popped) on read.
assert "req1" not in client.futures
def test_lookup_key_client_discard_clears_state():
"""discard() drops a completed lookup Future so it is not served stale."""
vllm_config = _make_vllm_config()
with patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
"worker.make_zmq_socket"
) as mock_make_socket:
client = worker.LookupKeyClient(vllm_config)
fake_socket = mock_make_socket.return_value
gate = threading.Event()
fake_socket.recv.side_effect = _gated_recv(gate, 9)
# Submit while gated so the call returns None and the Future stays in
# `futures` (unconsumed) once it resolves.
assert client.lookup("req2", 128, [], non_block=True) is None
gate.set()
deadline = time.monotonic() + 5.0
while time.monotonic() < deadline:
if client.futures["req2"].done():
break
time.sleep(0.005)
# discard() drops the completed result before any lookup consumes it.
client.discard("req2")
assert "req2" not in client.futures
# A fresh query re-submits rather than returning a stale value: hold the
# gate so the resubmitted lookup stays in flight.
gate.clear()
assert client.lookup("req2", 128, [], non_block=True) is None
gate.set() # release the executor so the worker thread can drain
def test_get_num_new_matched_tokens_async_defers_then_reports():
"""Async lookup returns (None, False) until ready, then the hit count."""
vllm_config = create_vllm_config(
kv_connector="MooncakeStoreConnector",
kv_role="kv_both",
kv_connector_extra_config={"lookup_async": True},
)
kv_cache_config = _make_kv_cache_config()
with (
set_current_vllm_config(vllm_config),
patch(
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
"scheduler.LookupKeyClient"
) as mock_client_cls,
):
sched = scheduler.MooncakeStoreScheduler(vllm_config, kv_cache_config)
assert sched.lookup_async is True
mock_client = mock_client_cls.return_value
block_size = sched._block_size
request = MagicMock()
request.request_id = "r1"
request.num_tokens = 4 * block_size
request.block_hashes = []
# Lookup not ready -> defer.
mock_client.lookup.return_value = None
assert sched.get_num_new_matched_tokens(request, 0) == (None, False)
assert "r1" not in sched.load_specs
# Lookup ready with a hit -> report need_to_allocate + async-load flag.
hit = 3 * block_size
mock_client.lookup.return_value = hit
need, load_async = sched.get_num_new_matched_tokens(request, 0)
assert need == hit
assert load_async == sched.load_async
assert sched.load_specs["r1"].kvpool_cached_tokens == hit
def test_protocol_tags_are_distinct_and_non_empty():
"""Protocol tags must be unique and non-empty to avoid collision."""
tags = {protocol.LOOKUP_MSG, protocol.RESET_MSG}
@@ -16,6 +16,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.scheduler impor
def _make_bare_scheduler() -> MooncakeStoreScheduler:
scheduler = object.__new__(MooncakeStoreScheduler)
scheduler.kv_role = "kv_both"
scheduler.lookup_async = False
scheduler._block_size = 16
scheduler.load_specs = {}
scheduler._preempted_req_ids = set()
@@ -405,7 +406,13 @@ class _StubLookupClient:
def __init__(self, hit_tokens: int) -> None:
self._hit_tokens = hit_tokens
def lookup(self, token_len: int, block_hashes: list[bytes]) -> int:
def lookup(
self,
req_id: str,
token_len: int,
block_hashes: list[bytes],
non_block: bool = False,
) -> int:
return self._hit_tokens
@@ -175,14 +175,17 @@ class _FakeModelConfig:
def _make_vllm_config(
*, extra_config: dict[str, object] | None = None
*,
extra_config: dict[str, object] | None = None,
rank: int = 0,
decode_context_parallel_size: int = 1,
) -> SimpleNamespace:
return SimpleNamespace(
model_config=_FakeModelConfig(),
parallel_config=SimpleNamespace(
pipeline_parallel_size=1,
rank=0,
decode_context_parallel_size=1,
rank=rank,
decode_context_parallel_size=decode_context_parallel_size,
prefill_context_parallel_size=1,
),
kv_transfer_config=_FakeKVTransferConfig(extra_config=extra_config),
@@ -231,13 +234,23 @@ def _install_fake_mooncake(monkeypatch, store_instance: MagicMock):
return FakeReplicateConfig
def _patch_worker_runtime(monkeypatch, *, local_ip: str = "10.0.0.7") -> None:
def _patch_worker_runtime(
monkeypatch,
*,
local_ip: str = "10.0.0.7",
tp_rank: int = 0,
tp_size: int = 1,
dcp_size: int = 1,
) -> None:
single_rank_group = SimpleNamespace(world_size=1, rank_in_group=0)
# DCP groups are contiguous splits of the TP group (see
# parallel_state.py), so dcp_rank == tp_rank % dcp_size.
dcp_group = SimpleNamespace(world_size=dcp_size, rank_in_group=tp_rank % dcp_size)
monkeypatch.setattr(worker, "get_mooncake_dp_engine_index", lambda _: 0)
monkeypatch.setattr(worker, "get_tensor_model_parallel_rank", lambda: 0)
monkeypatch.setattr(worker, "get_tensor_model_parallel_world_size", lambda: 1)
monkeypatch.setattr(worker, "get_tensor_model_parallel_rank", lambda: tp_rank)
monkeypatch.setattr(worker, "get_tensor_model_parallel_world_size", lambda: tp_size)
monkeypatch.setattr(worker, "get_pcp_group", lambda: single_rank_group)
monkeypatch.setattr(worker, "get_dcp_group", lambda: single_rank_group)
monkeypatch.setattr(worker, "get_dcp_group", lambda: dcp_group)
monkeypatch.setattr(worker, "get_ip", lambda: local_ip)
@@ -884,6 +897,66 @@ def test_requester_worker_init_builds_replicate_config_for_preferred_segment(
assert w.store_replicate_config.preferred_segment == "10.0.0.7:50053"
@pytest.mark.parametrize("dcp_size", [1, 4])
def test_worker_put_striding_covers_every_rank_get_namespace(
tmp_path, monkeypatch, dcp_size
):
"""Every key a rank GETs must have been PUT by some rank.
When num_kv_head < tp_size, ranks holding the same KV heads stripe
their PUTs across one shared key namespace. That dedup is only valid
when those ranks really share a namespace: with DCP > 1 each rank GETs
every key from its own ``@dcpN`` namespace, so striding must be
disabled.
"""
tp_size = 4
store = MagicMock()
store.setup.return_value = 0
_install_fake_mooncake(monkeypatch, store)
monkeypatch.setenv(
"MOONCAKE_CONFIG_PATH",
_write_mooncake_config(
tmp_path,
{
"metadata_server": "http://metadata/endpoint",
"protocol": "tcp",
"device_name": "",
"master_server_address": "10.0.0.7:50051",
},
),
)
# _FakeModelConfig has num_kv_head=1 < tp_size, which enables striding.
block_hashes = [f"hash-{i}".encode() for i in range(4)]
put_keys: set[str] = set()
get_keys_per_rank: dict[int, set[str]] = {}
for tp_rank in range(tp_size):
_patch_worker_runtime(
monkeypatch, tp_rank=tp_rank, tp_size=tp_size, dcp_size=dcp_size
)
w = worker.MooncakeStoreWorker(
_make_vllm_config(rank=tp_rank, decode_context_parallel_size=dcp_size),
_make_kv_cache_config(),
)
db = w.token_dbs[0]
token_len = len(block_hashes) * db.block_size
keys = [
key.to_string() for _, _, key in db.process_tokens(token_len, block_hashes)
]
assert len(keys) == len(block_hashes)
# PUT side: mirrors KVCacheStoreSendingThread's striding slice.
put_keys.update(keys[w.tp_rank % w.put_step :: w.put_step])
# GET side: KVCacheStoreRecvingThread fetches every key.
get_keys_per_rank[tp_rank] = set(keys)
for tp_rank, rank_keys in get_keys_per_rank.items():
missing = rank_keys - put_keys
assert not missing, (
f"tp_rank={tp_rank} would GET {len(missing)}/{len(rank_keys)} keys "
f"that no rank PUT (Mooncake OBJECT_NOT_FOUND): {sorted(missing)}"
)
# ---------------------------------------------------------------------------
# Helpers for register_kv_caches tests
# ---------------------------------------------------------------------------
+227
View File
@@ -0,0 +1,227 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for the Model Runner V2 Gumbel-max sampling kernel.
Accuracy: define a target categorical distribution as a non-negative int64
count tensor summing to N, turn it into logits (= log(count)), sample many
times with `gumbel_sample`, and check the empirical distribution matches.
The count tensor is deliberately heavy-tailed (one dominant token, the rest
~18 logits below). That tail is the sensitive part: the fp32 Gumbel noise must
reach ~18 to ever sample it. A flat distribution would keep every token within
a few logits of the top and would not exercise the noise tail at all.
"""
import math
import pytest
import torch
pytest.importorskip("triton")
if not torch.cuda.is_available():
pytest.skip("CUDA required for Gumbel sampler tests", allow_module_level=True)
from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
DEVICE = "cuda"
VOCAB_SIZE = 200_000
NUM_SAMPLES = 500_000
# Dominant token is exp(HEAD_LOG_GAP)x larger than the unit-count tail, so the
# tail sits ~HEAD_LOG_GAP logits below the top.
HEAD_LOG_GAP = 18.0
# 10-sigma band: a correct sampler effectively never trips it.
Z_TOLERANCE = 10.0
def _make_heavy_tailed_counts(seed: int = 1234) -> torch.Tensor:
"""Non-negative int64 counts of shape [VOCAB_SIZE]; target prob = counts/N."""
gen = torch.Generator(device=DEVICE).manual_seed(seed)
counts = torch.randint(
1, 4, (VOCAB_SIZE,), generator=gen, dtype=torch.int64, device=DEVICE
)
counts[0] = round(math.exp(HEAD_LOG_GAP)) # dominant token
return counts
def _counts_to_logits(counts: torch.Tensor) -> torch.Tensor:
# softmax(log(count)) == count / sum(count); count 0 -> logit -inf -> prob 0.
return counts.double().log().to(torch.float32)
def _sample(
logits_1d: torch.Tensor,
num_samples: int,
*,
use_fp64: bool = False,
temperature: float = 1.0,
) -> torch.Tensor:
"""Sample `num_samples` tokens from one logit vector.
Fixed seed with a distinct `pos` per sample gives independent draws; the
logits are broadcast with a 0-stride view to avoid materializing
[num_samples, vocab_size].
"""
vocab_size = logits_1d.shape[0]
logits = logits_1d.unsqueeze(0).expand(num_samples, vocab_size)
idx_mapping = torch.zeros(num_samples, dtype=torch.int32, device=DEVICE)
temp = torch.tensor([temperature], dtype=torch.float32, device=DEVICE)
seed = torch.tensor([0xABCD], dtype=torch.int64, device=DEVICE)
pos = torch.arange(num_samples, dtype=torch.int64, device=DEVICE)
return gumbel_sample(
logits,
idx_mapping,
temp,
seed,
pos,
apply_temperature=True,
use_fp64=use_fp64,
)
def _z_score(observed: int, expected: float, num_trials: int) -> float:
p = expected / num_trials
return (observed - expected) / math.sqrt(num_trials * p * (1 - p))
def _sample_histogram(
logits_1d: torch.Tensor, num_samples: int, *, chunk: int = 1_000_000
) -> torch.Tensor:
"""Histogram of `num_samples` draws, accumulated in chunks.
Chunking keeps the kernel's per-sample scratch ([chunk, num_blocks]) bounded
so a large sample count does not blow up memory.
"""
vocab_size = logits_1d.shape[0]
hist = torch.zeros(vocab_size, dtype=torch.float64, device=DEVICE)
for start in range(0, num_samples, chunk):
size = min(chunk, num_samples - start)
logits = logits_1d.unsqueeze(0).expand(size, vocab_size)
idx_mapping = torch.zeros(size, dtype=torch.int32, device=DEVICE)
temp = torch.tensor([1.0], dtype=torch.float32, device=DEVICE)
seed = torch.tensor([0xABCD], dtype=torch.int64, device=DEVICE)
pos = torch.arange(start, start + size, dtype=torch.int64, device=DEVICE)
out = gumbel_sample(
logits, idx_mapping, temp, seed, pos, apply_temperature=True
)
hist += torch.bincount(out, minlength=vocab_size).double()
return hist
# ----------------------------- Accuracy ------------------------------------
@pytest.mark.parametrize("use_fp64", [False, True])
def test_sampling_matches_target_distribution(use_fp64: bool):
counts = _make_heavy_tailed_counts()
total = counts.sum().item()
logits = _counts_to_logits(counts)
sampled = _sample(logits, NUM_SAMPLES, use_fp64=use_fp64)
assert sampled.min() >= 0 and sampled.max() < VOCAB_SIZE
# The dominant token (index 0) and the aggregate tail are the two
# statistically resolvable bins (individual tail tokens are far below the
# ~5/N detectability floor). The tail mass is small but well above noise,
# and it lives beyond the fp32 Gumbel cap -- the regime sensitive to noise
# precision -- so matching it is the meaningful check.
tail_prob = (total - counts[0].item()) / total
tail_count = (sampled != 0).sum().item()
z = _z_score(tail_count, NUM_SAMPLES * tail_prob, NUM_SAMPLES)
assert abs(z) < Z_TOLERANCE, (
f"sampled tail mass {tail_count / NUM_SAMPLES:.3e} != target "
f"{tail_prob:.3e} (z={z:.2f})"
)
def test_full_vocab_distribution_fidelity():
"""The sampled distribution matches the target across the WHOLE vocab.
A near-flat count tensor makes every one of the 200K bins individually
measurable. With ~20 samples/bin, a goodness-of-fit over all bins checks
that no part of the vocab is over- or under-represented (the heavy-tailed
test above only resolves head vs aggregate tail). Empirically the fp32
sampler is as faithful here as torch.multinomial; the residual error is the
multinomial sampling-noise floor, not the kernel.
"""
gen = torch.Generator(device=DEVICE).manual_seed(2024)
counts = torch.randint(
500, 1500, (VOCAB_SIZE,), generator=gen, dtype=torch.int64, device=DEVICE
)
total = counts.sum().item()
logits = _counts_to_logits(counts)
num_samples = 4_000_000
hist = _sample_histogram(logits, num_samples)
# Diversity: essentially every token must be reachable (no starved region).
coverage = (hist > 0).sum().item() / VOCAB_SIZE
assert coverage > 0.99, f"only {coverage:.4f} of the vocab was ever sampled"
# Goodness-of-fit across all bins (each has expected count >= ~10).
expected = (counts.double() / total) * num_samples
chi2 = (((hist - expected) ** 2) / expected).sum().item()
df = VOCAB_SIZE - 1
assert chi2 < df + 10 * math.sqrt(2 * df), f"chi2={chi2:.0f}, df={df}"
# ----------------------------- Edge cases ----------------------------------
def test_greedy_temperature_zero_returns_argmax():
"""temperature == 0 skips Gumbel noise and returns the exact argmax."""
torch.manual_seed(0)
num_reqs = 128
logits = torch.randn(num_reqs, VOCAB_SIZE, device=DEVICE, dtype=torch.float32)
idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=DEVICE)
temp = torch.zeros(num_reqs, dtype=torch.float32, device=DEVICE)
seed = torch.arange(num_reqs, dtype=torch.int64, device=DEVICE)
pos = torch.arange(num_reqs, dtype=torch.int64, device=DEVICE)
sampled = gumbel_sample(
logits, idx_mapping, temp, seed, pos, apply_temperature=True
)
assert torch.equal(sampled, logits.argmax(dim=-1))
def test_zero_count_tokens_are_never_sampled():
"""Count 0 -> -inf logit -> probability 0; must never be selected."""
counts = _make_heavy_tailed_counts(seed=7)
zeroed = torch.arange(1, VOCAB_SIZE, 2, device=DEVICE) # odd indices (not head)
counts[zeroed] = 0
logits = _counts_to_logits(counts)
sampled = _sample(logits, NUM_SAMPLES)
assert sampled.min() >= 0 and sampled.max() < VOCAB_SIZE
assert not torch.isin(sampled, zeroed).any(), "sampled a zero-probability token"
def test_single_nonzero_token_is_always_sampled():
"""A lone finite logit must win every draw, regardless of its index."""
counts = torch.zeros(VOCAB_SIZE, dtype=torch.int64, device=DEVICE)
counts[123_456] = 1000
logits = _counts_to_logits(counts)
sampled = _sample(logits, 10_000)
assert (sampled == 123_456).all()
@pytest.mark.parametrize("vocab_size", [1, 999, 1024, 4097])
def test_vocab_size_not_multiple_of_block(vocab_size: int):
"""Per-block tail masking for non-block-aligned vocab; all bins measurable."""
gen = torch.Generator(device=DEVICE).manual_seed(vocab_size)
counts = torch.randint(
20, 200, (vocab_size,), generator=gen, dtype=torch.int64, device=DEVICE
)
total = counts.sum().item()
logits = _counts_to_logits(counts)
num_samples = max(40 * vocab_size, 50_000)
sampled = _sample(logits, num_samples)
assert sampled.min() >= 0 and sampled.max() < vocab_size
observed = torch.bincount(sampled, minlength=vocab_size).double()
expected = (counts.double() / total) * num_samples
chi2 = (((observed - expected) ** 2) / expected).sum().item()
df = vocab_size - 1
if df >= 1:
assert chi2 < df + 10 * math.sqrt(2 * df), f"chi2={chi2:.1f}, df={df}"
+1 -1
View File
@@ -991,7 +991,7 @@ class VllmBackend:
},
payload_fn=lambda: json.dumps(
{
"model": self.vllm_config.model_config.model,
"model": getattr(self.vllm_config.model_config, "model", "unknown"),
"prefix": self.prefix,
"mode": str(cc.mode),
"backend": cc.backend,
+9
View File
@@ -302,6 +302,14 @@ class ParallelConfig:
Each entry must use `numactl --physcpubind` CPU-list syntax, for example
`"0-3"` or `"0,2,4-7"`.
"""
assigned_physical_gpu_ids: list[int] | None = None
"""Mapping from vLLM-local logical GPU IDs to physical GPU IDs.
For example, ``[2, 3]`` means logical GPU 0 maps to physical GPU 2,
and logical GPU 1 maps to physical GPU 3. Physical IDs are used only
at platform/topology boundaries such as NVML, NIC affinity, P2P
checks, and final CUDA device selection when needed. When None,
logical IDs map to visible device IDs in order."""
distributed_timeout_seconds: int | None = None
"""Timeout in seconds for distributed operations (e.g., init_process_group).
@@ -772,6 +780,7 @@ class ParallelConfig:
"numa_bind",
"numa_bind_nodes",
"numa_bind_cpus",
"assigned_physical_gpu_ids",
}
from vllm.config.utils import get_hash_factors, hash_factors
@@ -704,7 +704,14 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase):
self.num_experts = num_experts
self.cleanup()
gpus_per_node = torch.accelerator.device_count()
from vllm.platforms.interface import get_assigned_physical_gpu_ids
assigned_physical_gpu_ids = get_assigned_physical_gpu_ids()
gpus_per_node = (
len(assigned_physical_gpu_ids)
if assigned_physical_gpu_ids is not None
else torch.accelerator.device_count()
)
logger.debug(
"Making One-sided NVLink mapping: rank=%d, world size=%d",
self.rank,
@@ -320,13 +320,21 @@ def gpu_p2p_access_check(src: int, tgt: int) -> bool:
is_distributed = dist.is_initialized()
num_dev = current_platform.device_count()
cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES
if cuda_visible_devices is None:
cuda_visible_devices = ",".join(str(i) for i in range(num_dev))
from vllm.platforms.interface import get_assigned_physical_gpu_ids
assigned_physical_gpu_ids = get_assigned_physical_gpu_ids()
if assigned_physical_gpu_ids is not None:
# Key by the ordered list: the cache stores directed local-index
# pairs, so permutations of the same set are distinct mappings.
cache_key = ",".join(str(i) for i in assigned_physical_gpu_ids)
num_dev = len(assigned_physical_gpu_ids)
else:
num_dev = current_platform.device_count()
cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES
cache_key = cuda_visible_devices or ",".join(str(i) for i in range(num_dev))
path = os.path.join(
envs.VLLM_CACHE_ROOT, f"gpu_p2p_access_cache_for_{cuda_visible_devices}.json"
envs.VLLM_CACHE_ROOT, f"gpu_p2p_access_cache_for_{cache_key}.json"
)
os.makedirs(os.path.dirname(path), exist_ok=True)
from vllm.distributed.parallel_state import get_world_group
@@ -338,7 +346,15 @@ def gpu_p2p_access_check(src: int, tgt: int) -> bool:
# enter this block to calculate the cache
logger.info("generating GPU P2P access cache in %s", path)
cache: dict[str, bool] = {}
ids = list(range(num_dev))
# The probe subprocesses inherit this process's device-control env
# var, so they must be given visible ordinals, not physical IDs.
if assigned_physical_gpu_ids is not None:
ids = [
current_platform.logical_device_id_to_visible_device_id(local)
for local in range(num_dev)
]
else:
ids = list(range(num_dev))
# batch of all pairs of GPUs
batch_src, batch_tgt = zip(*list(product(ids, ids)))
# NOTE: we use `subprocess` rather than `multiprocessing` here
@@ -368,8 +384,11 @@ def gpu_p2p_access_check(src: int, tgt: int) -> bool:
) from e
with open(output_file.name, "rb") as f:
result = pickle.load(f)
# Cache entries must be keyed by local indices (0..N-1) because
# gpu_p2p_access_check() is called with local ranks.
id_to_local = {device_id: local for local, device_id in enumerate(ids)}
for _i, _j, r in zip(batch_src, batch_tgt, result):
cache[f"{_i}->{_j}"] = r
cache[f"{id_to_local[_i]}->{id_to_local[_j]}"] = r
with open(path, "w") as f:
json.dump(cache, f, indent=4)
if is_distributed:
@@ -34,7 +34,12 @@ def _can_p2p(rank: int, world_size: int) -> bool:
continue
if envs.VLLM_SKIP_P2P_CHECK:
logger.debug("Skipping P2P check and trusting the driver's P2P report.")
return torch.cuda.can_device_access_peer(rank, i)
# can_device_access_peer takes visible device ordinals, while
# rank and i are logical local IDs.
return torch.cuda.can_device_access_peer(
current_platform.logical_device_id_to_visible_device_id(rank),
current_platform.logical_device_id_to_visible_device_id(i),
)
if not gpu_p2p_access_check(rank, i):
return False
return True
@@ -126,13 +131,10 @@ class CustomAllreduce:
CUSTOM_ALL_REDUCE_MAX_SIZES[device_capability_str][world_size],
max_size,
)
cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES
if cuda_visible_devices:
device_ids = list(map(int, cuda_visible_devices.split(",")))
else:
device_ids = list(range(current_platform.device_count()))
physical_device_id = device_ids[device.index]
# device.index is a visible ordinal, not a logical local ID.
physical_device_id = current_platform.visible_device_id_to_physical_device_id(
device.index
)
tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu")
gather_list = [
torch.tensor([0], dtype=torch.int, device="cpu") for _ in range(world_size)
@@ -129,12 +129,10 @@ class QuickAllReduce:
assert isinstance(device, torch.device)
self.device = device
cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES
if cuda_visible_devices:
device_ids = list(map(int, cuda_visible_devices.split(",")))
else:
device_ids = list(range(current_platform.device_count()))
physical_device_id = device_ids[device.index]
# device.index is a visible ordinal, not a logical local ID.
physical_device_id = current_platform.visible_device_id_to_physical_device_id(
device.index
)
tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu")
gather_list = [
torch.tensor([0], dtype=torch.int, device="cpu")
@@ -840,7 +840,13 @@ class MessageQueue:
The MessageQueue instance for the calling process,
and a list of handles (only non-empty for the reader process).
"""
local_size = current_platform.device_count()
from vllm.platforms.interface import get_assigned_physical_gpu_ids
assigned_physical_gpu_ids = get_assigned_physical_gpu_ids()
if assigned_physical_gpu_ids is not None:
local_size = len(assigned_physical_gpu_ids)
else:
local_size = current_platform.device_count()
rank = dist.get_rank()
same_node = rank // local_size == reader_rank // local_size
buffer_io = MessageQueue(
@@ -482,10 +482,11 @@ def _init_lmcache_engine(
)
# Change current device.
num_gpus = torch.accelerator.device_count()
local_rank = parallel_config.rank % num_gpus
torch.accelerator.set_device_index(local_rank)
device = torch.device(f"cuda:{local_rank}")
from vllm.distributed.parallel_state import get_world_group
device_index = get_world_group().device_index
torch.accelerator.set_device_index(device_index)
device = torch.device(f"cuda:{device_index}")
metadata = LMCacheEngineMetadata(
model_config.model,
parallel_config.world_size,
@@ -176,7 +176,7 @@ class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA):
self,
request: Request,
num_computed_tokens: int,
) -> tuple[int, bool]:
) -> tuple[int | None, bool]:
assert self.connector_scheduler is not None
return self.connector_scheduler.get_num_new_matched_tokens(
request, num_computed_tokens
@@ -54,9 +54,9 @@ class MooncakeStoreScheduler:
):
assert vllm_config.kv_transfer_config is not None
self.kv_role = vllm_config.kv_transfer_config.kv_role
self.load_async = vllm_config.kv_transfer_config.kv_connector_extra_config.get(
"load_async", True
)
kvc_extra_config = vllm_config.kv_transfer_config.kv_connector_extra_config
self.load_async = kvc_extra_config.get("load_async", True)
self.lookup_async = kvc_extra_config.get("lookup_async", False)
self.client = LookupKeyClient(vllm_config)
# Align with the engine's own scheduler_block_size and hash_block_size.
@@ -75,14 +75,26 @@ class MooncakeStoreScheduler:
self,
request: Request,
num_computed_tokens: int,
) -> tuple[int, bool]:
"""Check for external KV cache hit."""
) -> tuple[int | None, bool]:
"""Check for external KV cache hit.
Returns ``(None, False)`` when an async lookup is still in flight,
signaling the scheduler to retry this request on a later step.
"""
# Look up against the full prefill range, not just the prompt.
token_len = request.num_tokens // self._block_size * self._block_size
if token_len < self._block_size:
return 0, False
num_external_hit_tokens = self.client.lookup(token_len, request.block_hashes)
num_external_hit_tokens = self.client.lookup(
request.request_id,
token_len,
request.block_hashes,
non_block=self.lookup_async,
)
if num_external_hit_tokens is None:
# Lookup not ready yet; scheduler will retry on a later step.
return None, False
if num_external_hit_tokens == request.num_tokens:
# Leave a sub-block tail uncomputed for sampling, on a block
@@ -158,6 +170,7 @@ class MooncakeStoreScheduler:
force_skip_save = self.kv_role == "kv_consumer"
for finished_req_id in scheduler_output.finished_req_ids:
self.client.discard(finished_req_id)
self.load_specs.pop(finished_req_id, None)
self._request_trackers.pop(finished_req_id, None)
self._unfinished_requests.pop(finished_req_id, None)
@@ -19,6 +19,7 @@ import threading
import time
from collections import defaultdict
from collections.abc import Callable
from concurrent.futures import Future, ThreadPoolExecutor
from dataclasses import dataclass
from typing import Any, Literal, TypeVar
@@ -971,7 +972,13 @@ class MooncakeStoreWorker:
else:
self.num_kv_head = model_config.get_total_num_kv_heads()
if self.num_kv_head < self.tp_size:
if self.num_kv_head < self.tp_size and self.dcp_size <= 1:
# Dedup: TP ranks holding the same KV heads stripe PUTs across
# one shared key namespace. DCP splits the TP group, so with
# DCP>1 those ranks have different `@dcpN` namespaces and
# striping would leave keys unwritten (OBJECT_NOT_FOUND on
# GET). PCP is outer to TP (pcp_rank is constant within a TP
# group), so it needs no guard.
self.put_step = self.tp_size // self.num_kv_head
self.head_or_tp_rank = self.tp_rank // self.put_step
else:
@@ -1560,7 +1567,13 @@ class LookupKeyClient:
bind=False,
)
def lookup(self, token_len: int, block_hashes: list[BlockHash]) -> int:
# Async lookup support
self.executor = ThreadPoolExecutor(
max_workers=1, thread_name_prefix="MooncakeLookupClient"
)
self.futures: dict[str, Future[int]] = {}
def _lookup(self, token_len: int, block_hashes: list[BlockHash]) -> int:
hash_strs = [h.hex() for h in block_hashes]
hash_frames = self.encoder.encode(hash_strs)
token_len_bytes = token_len.to_bytes(4, byteorder="big")
@@ -1570,7 +1583,36 @@ class LookupKeyClient:
result = int.from_bytes(resp, "big")
return result
def reset(self) -> bool:
def lookup(
self,
req_id: str,
token_len: int,
block_hashes: list[BlockHash],
non_block: bool = False,
) -> int | None:
"""If non_block is True, will return None until the result is ready,
so the caller retries on a later step."""
future = self.futures.get(req_id)
if future is None:
future = self.executor.submit(self._lookup, token_len, list(block_hashes))
self.futures[req_id] = future
if non_block and not future.done():
return None
try:
return future.result()
except Exception as e:
logger.error("Async Mooncake lookup failed for %s: %s", req_id, e)
return 0
finally:
del self.futures[req_id]
def discard(self, req_id: str) -> None:
"""Drop any cached/in-flight lookup for ``req_id`` (e.g. on abort)."""
future = self.futures.pop(req_id, None)
if future is not None:
future.cancel()
def _reset(self) -> bool:
"""Trigger ``store.remove_all(force=True)`` on worker rank 0.
Ordering assumption: caller MUST ensure no in-flight Mooncake
@@ -1582,7 +1624,11 @@ class LookupKeyClient:
resp = self.socket.recv()
return bytes(resp) == RESP_OK
def reset(self) -> bool:
return self.executor.submit(self._reset).result()
def close(self):
self.executor.shutdown(wait=False, cancel_futures=True)
self.socket.close(linger=0)
@@ -841,8 +841,106 @@ class NixlBaseConnectorWorker:
# Forwarding a real layer name rather than a synthetic key
self.register_kv_caches({first_layer: kv_cache})
def _register_packed_kv_cache(
self,
storage: torch.UntypedStorage,
) -> None:
"""Register a packed KV cache as a single NIXL region.
The packed allocation interleaves all layers per block, so each
block_stride-byte chunk is one logical block. We register 1
NIXL region and create 1 descriptor per block.
"""
self.transfer_topo = TransferTopology(
tp_rank=self.tp_rank,
tp_size=self.world_size,
block_size=self.block_size,
engine_id=self.engine_id,
is_mla=self.use_mla,
total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
attn_backends=self.attn_backends,
tensor_shape=None,
is_mamba=self._has_mamba,
)
self.compat_hash = compute_nixl_compatibility_hash(
self.vllm_config,
self.backend_name,
self.transfer_topo.cross_layers_blocks,
)
total_size = storage.nbytes()
block_stride = total_size // self.num_blocks
base_addr = storage.data_ptr()
device_id = storage.device.index
assert device_id is not None
logger.info(
"Registering packed KV cache: total_size=%s, block_stride=%s, "
"num_blocks=%s, num_regions=1",
total_size,
block_stride,
self.num_blocks,
)
self.device_id = device_id
caches_data = [(base_addr, total_size, self.device_id, "")]
self.block_len_per_layer = [block_stride]
self.num_regions = 1
self.num_descs = self.num_blocks
self.kv_caches_base_addr[self.engine_id][self.tp_rank] = [base_addr]
descs = self.nixl_wrapper.get_reg_descs(caches_data, self.nixl_memory_type)
self.nixl_wrapper.register_memory(descs, backends=self.nixl_backends)
self._registered_descs.append(descs)
self.dst_num_blocks[self.engine_id] = self.num_blocks
self.src_xfer_handles_by_block_size[self.block_size], (self.src_blocks_data) = (
self.register_local_xfer_handler(self.block_size)
)
agent_metadata = NixlAgentMetadata(
engine_id=self.engine_id,
agent_metadata=self.nixl_wrapper.get_agent_metadata(),
device_id=self.device_id,
kv_caches_base_addr=(
self.kv_caches_base_addr[self.engine_id][self.tp_rank]
),
num_blocks=self.num_blocks,
block_lens=self.block_len_per_layer,
kv_cache_layout=self.kv_cache_layout,
block_size=self.block_size,
ssm_sizes=self._mamba_ssm_size,
attn_backend_name=self.backend_name,
physical_blocks_per_logical_kv_block=(
self._physical_blocks_per_logical_kv_block
),
)
assert self.compat_hash is not None
encoder = msgspec.msgpack.Encoder()
self.xfer_handshake_metadata = NixlHandshakePayload(
compatibility_hash=self.compat_hash,
agent_metadata_bytes=encoder.encode(agent_metadata),
)
def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
"""Register the KV Cache data in nixl."""
# Detect packed allocation: all tensors are strided views into the
# same backing storage (different data_ptr but same storage).
# This happens with DSv4-style contiguous per-block packing.
if len(kv_caches) > 1 and not self._has_mamba:
storage = next(iter(kv_caches.values())).untyped_storage()
storage_ptrs = {
cache.untyped_storage().data_ptr() for cache in kv_caches.values()
}
data_ptrs = {cache.data_ptr() for cache in kv_caches.values()}
if len(storage_ptrs) == 1 and len(data_ptrs) > 1:
self._register_packed_kv_cache(storage)
self.device_kv_caches = kv_caches
return
self.transfer_topo = TransferTopology(
tp_rank=self.tp_rank,
tp_size=self.world_size,
@@ -72,18 +72,22 @@ class OffloadingConnectorWorker:
if isinstance(layer_kv_cache_spec, AttentionSpec):
layer_kv_cache = kv_caches[layer_name]
assert isinstance(layer_kv_cache, torch.Tensor)
assert layer_kv_cache.storage_offset() == 0
storage = layer_kv_cache.untyped_storage()
page = layer_kv_cache_spec.page_size_bytes
elem_size = layer_kv_cache.element_size()
byte_offset = layer_kv_cache.storage_offset() * elem_size
block_stride_bytes = layer_kv_cache.stride(0) * elem_size
tensors_per_block[layer_name] = (
torch.tensor(
[],
dtype=torch.int8,
device=layer_kv_cache.device,
)
.set_(storage)
.view(num_blocks, page),
).set_(
layer_kv_cache.untyped_storage(),
byte_offset,
(num_blocks, page),
(block_stride_bytes, 1),
),
)
page_size_bytes[layer_name] = layer_kv_cache_spec.page_size_bytes
unpadded_page_size_bytes[layer_name] = (
+24 -4
View File
@@ -392,6 +392,14 @@ class GroupCoordinator:
self.rank = torch.distributed.get_rank()
self.local_rank = local_rank
self.device_index: int
if _WORLD is not None:
self.device_index = _WORLD.device_index
else:
assert local_rank >= 0, (
"local_rank must be provided when creating the world group"
)
self.device_index = local_rank
self_device_group = None
self_cpu_group = None
@@ -442,11 +450,18 @@ class GroupCoordinator:
from vllm.platforms import current_platform
if current_platform.is_cuda_alike():
self.device = torch.device(f"cuda:{local_rank}")
visible_device_index = (
current_platform.logical_device_id_to_visible_device_id(
self.device_index
)
)
self.device = torch.device(f"cuda:{visible_device_index}")
elif current_platform.is_xpu():
self.device = torch.device(f"xpu:{local_rank}")
self.device = torch.device(f"xpu:{self.device_index}")
elif current_platform.is_out_of_tree():
self.device = torch.device(f"{current_platform.device_name}:{local_rank}")
self.device = torch.device(
f"{current_platform.device_name}:{self.device_index}"
)
else:
self.device = torch.device("cpu")
@@ -1438,7 +1453,12 @@ def _init_process_group_for_split_group(
"""
if torch.accelerator.is_available() and backend != "gloo":
init_backend = "cpu:gloo,cuda:nccl"
device_id: torch.device | None = torch.device(f"cuda:{local_rank}")
from vllm.platforms import current_platform
visible_device_index = current_platform.logical_device_id_to_visible_device_id(
local_rank
)
device_id: torch.device | None = torch.device(f"cuda:{visible_device_index}")
else:
init_backend = "gloo"
device_id = None
+19 -3
View File
@@ -86,6 +86,15 @@ class StatelessGroupCoordinator(GroupCoordinator):
self.rank = global_rank
self.local_rank = local_rank
from vllm.distributed.parallel_state import _WORLD
if _WORLD is not None:
self.device_index = _WORLD.device_index
else:
assert local_rank >= 0, (
"local_rank must be provided when creating the world group"
)
self.device_index = local_rank
self_device_group = None
self_cpu_group = None
@@ -152,11 +161,18 @@ class StatelessGroupCoordinator(GroupCoordinator):
self.tcp_store_group = self_tcp_store_group
if current_platform.is_cuda_alike():
self.device = torch.device(f"cuda:{local_rank}")
visible_device_index = (
current_platform.logical_device_id_to_visible_device_id(
self.device_index
)
)
self.device = torch.device(f"cuda:{visible_device_index}")
elif current_platform.is_xpu():
self.device = torch.device(f"xpu:{local_rank}")
self.device = torch.device(f"xpu:{self.device_index}")
elif current_platform.is_out_of_tree():
self.device = torch.device(f"{current_platform.device_name}:{local_rank}")
self.device = torch.device(
f"{current_platform.device_name}:{self.device_index}"
)
else:
self.device = torch.device("cpu")
+58
View File
@@ -6,6 +6,7 @@ import copy
import dataclasses
import functools
import json
import os
import sys
from collections.abc import Callable
from dataclasses import MISSING, asdict, dataclass, fields, is_dataclass
@@ -465,6 +466,7 @@ class EngineArgs:
numa_bind: bool = ParallelConfig.numa_bind
numa_bind_nodes: list[int] | None = ParallelConfig.numa_bind_nodes
numa_bind_cpus: list[str] | None = ParallelConfig.numa_bind_cpus
device_ids: list[int | str] | None = None
tensor_parallel_size: int = ParallelConfig.tensor_parallel_size
prefill_context_parallel_size: int = ParallelConfig.prefill_context_parallel_size
decode_context_parallel_size: int = ParallelConfig.decode_context_parallel_size
@@ -979,6 +981,20 @@ class EngineArgs:
parallel_group.add_argument(
"--numa-bind-cpus", **parallel_kwargs["numa_bind_cpus"]
)
parallel_group.add_argument(
"--device-ids",
type=lambda s: [
int(device_id) if device_id.isdigit() else device_id
for device_id in (part.strip() for part in s.split(","))
],
default=None,
help="Comma-separated physical GPU device IDs or UUIDs to use "
'(e.g. --device-ids "2,3,5,7"). Avoids setting '
"CUDA_VISIBLE_DEVICES, preserving full GPU topology "
"visibility for GPU-NIC affinity and DeepGEMM. "
"Note: has no effect with Ray executors; use Ray "
"placement groups for GPU selection instead.",
)
parallel_group.add_argument(
"--tensor-parallel-size", "-tp", **parallel_kwargs["tensor_parallel_size"]
)
@@ -1716,6 +1732,47 @@ class EngineArgs:
)
return SpeculativeConfig(**self.speculative_config)
def _resolve_device_ids(self) -> list[int] | None:
if not self.device_ids:
return None
if self.distributed_executor_backend == "ray":
logger.warning(
"--device-ids has no effect when using the Ray executor. "
"Use Ray placement groups for GPU selection instead."
)
ids = self.device_ids
if len(set(ids)) != len(ids):
raise ValueError(f"--device-ids must not contain duplicates: {ids}")
if all(isinstance(i, str) for i in ids):
return [
current_platform.device_control_id_to_physical_device_id(i)
for i in cast(list[str], ids)
]
if any(isinstance(i, str) for i in ids):
raise ValueError("--device-ids must not mix integer IDs and UUIDs")
int_ids = cast(list[int], ids)
# Compose with CUDA_VISIBLE_DEVICES: if CVD is set, treat
# --device-ids values as indices into the CVD-visible set.
cvd = getattr(
envs,
current_platform.device_control_env_var,
os.environ.get(current_platform.device_control_env_var),
)
if cvd:
cvd_ids = [
current_platform.device_control_id_to_physical_device_id(x)
for x in cvd.split(",")
]
for i in int_ids:
if i >= len(cvd_ids):
raise ValueError(
f"--device-ids index {i} is out of range for "
f"{current_platform.device_control_env_var}"
f"={cvd} ({len(cvd_ids)} devices visible)"
)
return [cvd_ids[i] for i in int_ids]
return int_ids
def create_diffusion_config(self) -> DiffusionConfig | None:
if self.diffusion_config is None:
return None
@@ -2029,6 +2086,7 @@ class EngineArgs:
cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
_api_process_count=self._api_process_count,
_api_process_rank=self._api_process_rank,
assigned_physical_gpu_ids=self._resolve_device_ids(),
numa_bind=self.numa_bind,
numa_bind_nodes=self.numa_bind_nodes,
numa_bind_cpus=self.numa_bind_cpus,
+73 -6
View File
@@ -12,6 +12,7 @@ import uuid
from collections.abc import AsyncGenerator
from typing import TYPE_CHECKING, Any
import jinja2
from fastapi import Request
from vllm.engine.protocol import EngineClient
@@ -99,6 +100,36 @@ class AnthropicServingMessages(OpenAIServingChat):
"length": "max_tokens",
"tool_calls": "tool_use",
}
self._merge_inline_system = self._detect_merge_inline_system(chat_template)
@staticmethod
def _detect_merge_inline_system(chat_template: str | None) -> bool:
"""Auto-detect whether the chat template requires system-first ordering.
Renders a [system, user, system, user] conversation against the
template; if it raises (e.g. Qwen's ``loop.first`` guard), the
model needs inline system messages merged into the leading block.
"""
if not chat_template:
return True
try:
env = jinja2.sandbox.ImmutableSandboxedEnvironment(
trim_blocks=True,
lstrip_blocks=True,
extensions=[jinja2.ext.loopcontrols],
)
env.from_string(chat_template).render(
messages=[
{"role": "system", "content": "t"},
{"role": "user", "content": "t"},
{"role": "system", "content": "t"},
{"role": "user", "content": "t"},
],
add_generation_prompt=False,
)
return False
except jinja2.TemplateError:
return True
@staticmethod
def _convert_image_source_to_url(source: dict[str, Any]) -> str:
@@ -123,13 +154,24 @@ class AnthropicServingMessages(OpenAIServingChat):
@classmethod
def _convert_anthropic_to_openai_request(
cls, anthropic_request: AnthropicMessagesRequest | AnthropicCountTokensRequest
cls,
anthropic_request: AnthropicMessagesRequest | AnthropicCountTokensRequest,
*,
merge_inline_system: bool = False,
) -> ChatCompletionRequest:
"""Convert Anthropic message format to OpenAI format"""
openai_messages: list[dict[str, Any]] = []
cls._convert_system_message(anthropic_request, openai_messages)
cls._convert_messages(anthropic_request.messages, openai_messages)
cls._convert_system_message(
anthropic_request,
openai_messages,
merge_inline_system=merge_inline_system,
)
cls._convert_messages(
anthropic_request.messages,
openai_messages,
merge_inline_system=merge_inline_system,
)
req = cls._build_base_request(anthropic_request, openai_messages)
cls._handle_streaming_options(req, anthropic_request)
cls._handle_output_config(req, anthropic_request)
@@ -142,6 +184,8 @@ class AnthropicServingMessages(OpenAIServingChat):
cls,
anthropic_request: AnthropicMessagesRequest | AnthropicCountTokensRequest,
openai_messages: list[dict[str, Any]],
*,
merge_inline_system: bool = False,
) -> None:
"""Convert Anthropic system message to OpenAI format"""
system_parts: list[str] = []
@@ -159,6 +203,17 @@ class AnthropicServingMessages(OpenAIServingChat):
continue
system_parts.append(block.text)
# When the template requires system-first ordering, extract inline
# system messages from the messages array and merge them into the
# top-level block so the template doesn't reject them.
if merge_inline_system:
for msg in anthropic_request.messages:
if msg.role != "system":
continue
text = cls._extract_system_text(msg)
if text:
system_parts.append(text)
if system_parts:
openai_messages.append({"role": "system", "content": "".join(system_parts)})
@@ -180,7 +235,11 @@ class AnthropicServingMessages(OpenAIServingChat):
@classmethod
def _convert_messages(
cls, messages: list, openai_messages: list[dict[str, Any]]
cls,
messages: list,
openai_messages: list[dict[str, Any]],
*,
merge_inline_system: bool = False,
) -> None:
"""Convert Anthropic messages to OpenAI format"""
for msg in messages:
@@ -190,6 +249,8 @@ class AnthropicServingMessages(OpenAIServingChat):
# doesn't strip billing headers and may produce messages with
# no "content" key.
if msg.role == "system":
if merge_inline_system:
continue # already merged into top-level by _convert_system_message
text = cls._extract_system_text(msg)
if text:
openai_messages.append({"role": "system", "content": text})
@@ -497,7 +558,10 @@ class AnthropicServingMessages(OpenAIServingChat):
"""
if logger.isEnabledFor(logging.DEBUG):
logger.debug("Received messages request %s", request.model_dump_json())
chat_req = self._convert_anthropic_to_openai_request(request)
chat_req = self._convert_anthropic_to_openai_request(
request,
merge_inline_system=self._merge_inline_system,
)
if logger.isEnabledFor(logging.DEBUG):
logger.debug("Convert to OpenAI request %s", chat_req.model_dump_json())
generator = await self.create_chat_completion(chat_req, raw_request)
@@ -905,7 +969,10 @@ class AnthropicServingMessages(OpenAIServingChat):
raw_request: Request | None = None,
) -> AnthropicCountTokensResponse | ErrorResponse:
"""Implements Anthropic's messages.count_tokens endpoint."""
chat_req = self._convert_anthropic_to_openai_request(request)
chat_req = self._convert_anthropic_to_openai_request(
request,
merge_inline_system=self._merge_inline_system,
)
result = await self.render_chat_request(chat_req)
if isinstance(result, ErrorResponse):
return result
-6
View File
@@ -898,12 +898,6 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin, OfflineInferenceMixin):
def finish_weight_update(self) -> None:
"""Finish the current weight update."""
self.llm_engine.collective_rpc("finish_weight_update")
# Invalidate cached state computed with the old weights so it isn't
# reused for subsequent requests:
# - prefix cache: KV blocks computed with the old weights
# - encoder cache: multimodal embeddings keyed only by mm_hash
self.llm_engine.reset_prefix_cache()
self.llm_engine.reset_encoder_cache()
def __repr__(self) -> str:
"""Return a transformers-style hierarchical view of the model."""

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